




"""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 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase  ,      _UpperCAmelCase	):




	"""simple docstring"""






	if isinstance(lowerCamelCase__  ,      torch.Tensor	):
		return image
	elif isinstance(lowerCamelCase__  ,      PIL.Image.Image	):
		A_  :				Optional[Any]	   =		[image]

	if isinstance(image[0]  ,      PIL.Image.Image	):
		A_  :				List[Any]	   =		[np.array(i.resize((w, h)  ,      resample=PIL_INTERPOLATION['''lanczos''']	)	)[None, :] for i in image]
		A_  :				Union[str, Any]	   =		np.concatenate(lowerCamelCase__  ,      axis=0	)
		A_  :				Optional[int]	   =		np.array(lowerCamelCase__	).astype(np.floataa	) / 255.0
		A_  :				int	   =		image.transpose(0  ,      3  ,      1  ,      2	)
		A_  :				str	   =		2.0 * image - 1.0
		A_  :				List[str]	   =		torch.from_numpy(lowerCamelCase__	)
	elif isinstance(image[0]  ,      torch.Tensor	):
		A_  :				Union[str, Any]	   =		torch.cat(lowerCamelCase__  ,      dim=0	)
	return image

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase  ,      _UpperCAmelCase  ,      _UpperCAmelCase=0.9995	):




	"""simple docstring"""






	if not isinstance(lowerCamelCase__  ,      np.ndarray	):
		A_  :				Union[str, Any]	   =		True
		A_  :				str	   =		va.device
		A_  :				str	   =		va.cpu().numpy()
		A_  :				int	   =		va.cpu().numpy()

	A_  :				Optional[Any]	   =		np.sum(va * va / (np.linalg.norm(lowerCamelCase__	) * np.linalg.norm(lowerCamelCase__	))	)
	if np.abs(lowerCamelCase__	) > DOT_THRESHOLD:
		A_  :				Dict	   =		(1 - t) * va + t * va
	else:
		A_  :				Tuple	   =		np.arccos(lowerCamelCase__	)
		A_  :				Tuple	   =		np.sin(lowerCamelCase__	)
		A_  :				Optional[Any]	   =		theta_a * t
		A_  :				Union[str, Any]	   =		np.sin(lowerCamelCase__	)
		A_  :				Dict	   =		np.sin(theta_a - theta_t	) / sin_theta_a
		A_  :				Any	   =		sin_theta_t / sin_theta_a
		A_  :				Union[str, Any]	   =		sa * va + sa * va

	if inputs_are_torch:
		A_  :				Tuple	   =		torch.from_numpy(lowerCamelCase__	).to(lowerCamelCase__	)

	return va

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase	):




	"""simple docstring"""






	A_  :				Optional[Any]	   =		F.normalize(lowerCamelCase__  ,      dim=-1	)
	A_  :				str	   =		F.normalize(lowerCamelCase__  ,      dim=-1	)
	return (x - y).norm(dim=-1	).div(2	).arcsin().pow(2	).mul(2	)

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase	):




	"""simple docstring"""






	for param in model.parameters():
		A_  :				Optional[Any]	   =		value

class        lowercase      (    _UpperCamelCase):







				def __init__(							self   :					Any      ,     _lowerCamelCase   :					AutoencoderKL      ,     _lowerCamelCase   :					CLIPTextModel      ,     _lowerCamelCase   :					CLIPModel      ,     _lowerCamelCase   :					CLIPTokenizer      ,     _lowerCamelCase   :					UNetaDConditionModel      ,     _lowerCamelCase   :					Union[PNDMScheduler, LMSDiscreteScheduler, DDIMScheduler, DPMSolverMultistepScheduler]      ,     _lowerCamelCase   :					CLIPFeatureExtractor      ,     _lowerCamelCase   :					Any=None      ,     _lowerCamelCase   :					Tuple=None      ,     _lowerCamelCase   :					Tuple=None      ,     ):



					"""simple docstring"""

					super().__init__()
					self.register_modules(
					    vae=_lowerCamelCase      ,     text_encoder=_lowerCamelCase      ,     clip_model=_lowerCamelCase      ,     tokenizer=_lowerCamelCase      ,     unet=_lowerCamelCase      ,     scheduler=_lowerCamelCase      ,     feature_extractor=_lowerCamelCase      ,     coca_model=_lowerCamelCase      ,     coca_tokenizer=_lowerCamelCase      ,     coca_transform=_lowerCamelCase      ,     )
					A_  :				Dict	   =		(
					    feature_extractor.size
					    if isinstance(feature_extractor.size      ,     _lowerCamelCase       )
					    else feature_extractor.size["shortest_edge"]
					)
					A_  :				List[Any]	   =		transforms.Normalize(mean=feature_extractor.image_mean      ,     std=feature_extractor.image_std       )
					set_requires_grad(self.text_encoder      ,     _lowerCamelCase       )
					set_requires_grad(self.clip_model      ,     _lowerCamelCase       )







				def 					a_  (							self   :					int      ,     _lowerCamelCase   :					Optional[Union[str, int]] = "auto"       ):



					"""simple docstring"""

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







				def 					a_  (							self   :					List[Any]       ):



					"""simple docstring"""

					self.enable_attention_slicing(_lowerCamelCase       )







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

					set_requires_grad(self.vae      ,     _lowerCamelCase       )







				def 					a_  (							self   :					Tuple       ):



					"""simple docstring"""

					set_requires_grad(self.vae      ,     _lowerCamelCase       )







				def 					a_  (							self   :					List[str]       ):



					"""simple docstring"""

					set_requires_grad(self.unet      ,     _lowerCamelCase       )







				def 					a_  (							self   :					int       ):



					"""simple docstring"""

					set_requires_grad(self.unet      ,     _lowerCamelCase       )







				def 					a_  (							self   :					Tuple      ,     _lowerCamelCase   :					Optional[int]      ,     _lowerCamelCase   :					List[Any]      ,     _lowerCamelCase   :					Optional[int]       ):



					"""simple docstring"""

					A_  :				Optional[int]	   =		min(int(num_inference_steps * strength       )      ,     _lowerCamelCase       )

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

					return timesteps, num_inference_steps - t_start







				def 					a_  (							self   :					Dict      ,     _lowerCamelCase   :					Optional[int]      ,     _lowerCamelCase   :					Tuple      ,     _lowerCamelCase   :					Dict      ,     _lowerCamelCase   :					Any      ,     _lowerCamelCase   :					Dict      ,     _lowerCamelCase   :					str=None       ):



					"""simple docstring"""

					if not isinstance(_lowerCamelCase      ,     torch.Tensor       ):
						raise ValueError(F"""`image` has to be of type `torch.Tensor` but is {type(_lowerCamelCase       )}"""       )

					A_  :				Optional[int]	   =		image.to(device=_lowerCamelCase      ,     dtype=_lowerCamelCase       )

					if isinstance(_lowerCamelCase      ,     _lowerCamelCase       ):
						A_  :				List[str]	   =		[
						    self.vae.encode(image[i : i + 1]       ).latent_dist.sample(generator[i]       ) for i in range(_lowerCamelCase       )
						]
						A_  :				Any	   =		torch.cat(_lowerCamelCase      ,     dim=0       )
					else:
						A_  :				Optional[int]	   =		self.vae.encode(_lowerCamelCase       ).latent_dist.sample(_lowerCamelCase       )

					# Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
					A_  :				str	   =		0.18215 * init_latents
					A_  :				int	   =		init_latents.repeat_interleave(_lowerCamelCase      ,     dim=0       )

					A_  :				Dict	   =		randn_tensor(init_latents.shape      ,     generator=_lowerCamelCase      ,     device=_lowerCamelCase      ,     dtype=_lowerCamelCase       )

					# get latents
					A_  :				Tuple	   =		self.scheduler.add_noise(_lowerCamelCase      ,     _lowerCamelCase      ,     _lowerCamelCase       )
					A_  :				Any	   =		init_latents

					return latents







				def 					a_  (							self   :					int      ,     _lowerCamelCase   :					Optional[int]       ):



					"""simple docstring"""

					A_  :				Union[str, Any]	   =		self.coca_transform(_lowerCamelCase       ).unsqueeze(0       )
					with torch.no_grad(), torch.cuda.amp.autocast():
						A_  :				Dict	   =		self.coca_model.generate(transformed_image.to(device=self.device      ,     dtype=self.coca_model.dtype       )       )
					A_  :				List[str]	   =		self.coca_tokenizer.decode(generated[0].cpu().numpy()       )
					return generated.split('''<end_of_text>'''       )[0].replace('''<start_of_text>'''      ,     ''''''       ).rstrip(''' .,'''       )







				def 					a_  (							self   :					List[Any]      ,     _lowerCamelCase   :					int      ,     _lowerCamelCase   :					str       ):



					"""simple docstring"""

					A_  :				List[str]	   =		self.feature_extractor.preprocess(_lowerCamelCase       )
					A_  :				Optional[Any]	   =		torch.from_numpy(clip_image_input['''pixel_values'''][0]       ).unsqueeze(0       ).to(self.device       ).half()
					A_  :				int	   =		self.clip_model.get_image_features(_lowerCamelCase       )
					A_  :				Optional[int]	   =		image_embeddings_clip / image_embeddings_clip.norm(p=2      ,     dim=-1      ,     keepdim=_lowerCamelCase       )
					A_  :				List[str]	   =		image_embeddings_clip.repeat_interleave(_lowerCamelCase      ,     dim=0       )
					return image_embeddings_clip







				@torch.enable_grad()
				def 					a_  (							self   :					Optional[Any]      ,     _lowerCamelCase   :					List[Any]      ,     _lowerCamelCase   :					List[Any]      ,     _lowerCamelCase   :					Dict      ,     _lowerCamelCase   :					Optional[int]      ,     _lowerCamelCase   :					Dict      ,     _lowerCamelCase   :					Optional[Any]      ,     _lowerCamelCase   :					Any      ,     ):



					"""simple docstring"""

					A_  :				Dict	   =		latents.detach().requires_grad_()

					A_  :				Union[str, Any]	   =		self.scheduler.scale_model_input(_lowerCamelCase      ,     _lowerCamelCase       )

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

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

						A_  :				int	   =		torch.sqrt(_lowerCamelCase       )
						A_  :				Optional[int]	   =		pred_original_sample * (fac) + latents * (1 - fac)
					elif isinstance(self.scheduler      ,     _lowerCamelCase       ):
						A_  :				str	   =		self.scheduler.sigmas[index]
						A_  :				Dict	   =		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
					A_  :				Union[str, Any]	   =		1 / 0.18215 * sample
					A_  :				Dict	   =		self.vae.decode(_lowerCamelCase       ).sample
					A_  :				Union[str, Any]	   =		(image / 2 + 0.5).clamp(0      ,     1       )

					A_  :				Any	   =		transforms.Resize(self.feature_extractor_size       )(_lowerCamelCase       )
					A_  :				Optional[Any]	   =		self.normalize(_lowerCamelCase       ).to(latents.dtype       )

					A_  :				List[str]	   =		self.clip_model.get_image_features(_lowerCamelCase       )
					A_  :				Any	   =		image_embeddings_clip / image_embeddings_clip.norm(p=2      ,     dim=-1      ,     keepdim=_lowerCamelCase       )

					A_  :				Any	   =		spherical_dist_loss(_lowerCamelCase      ,     _lowerCamelCase       ).mean() * clip_guidance_scale

					A_  :				Tuple	   =		-torch.autograd.grad(_lowerCamelCase      ,     _lowerCamelCase       )[0]

					if isinstance(self.scheduler      ,     _lowerCamelCase       ):
						A_  :				Optional[int]	   =		latents.detach() + grads * (sigma**2)
						A_  :				str	   =		noise_pred_original
					else:
						A_  :				Union[str, Any]	   =		noise_pred_original - torch.sqrt(_lowerCamelCase       ) * grads
					return noise_pred, latents







				@torch.no_grad()
				def __call__(							self   :					List[str]      ,     _lowerCamelCase   :					Union[torch.FloatTensor, PIL.Image.Image]      ,     _lowerCamelCase   :					Union[torch.FloatTensor, PIL.Image.Image]      ,     _lowerCamelCase   :					Optional[str] = None      ,     _lowerCamelCase   :					Optional[str] = None      ,     _lowerCamelCase   :					Optional[int] = 5_12      ,     _lowerCamelCase   :					Optional[int] = 5_12      ,     _lowerCamelCase   :					float = 0.6      ,     _lowerCamelCase   :					Optional[int] = 50      ,     _lowerCamelCase   :					Optional[float] = 7.5      ,     _lowerCamelCase   :					Optional[int] = 1      ,     _lowerCamelCase   :					float = 0.0      ,     _lowerCamelCase   :					Optional[float] = 1_00      ,     _lowerCamelCase   :					Optional[torch.Generator] = None      ,     _lowerCamelCase   :					Optional[str] = "pil"      ,     _lowerCamelCase   :					bool = True      ,     _lowerCamelCase   :					float = 0.8      ,     _lowerCamelCase   :					float = 0.1      ,     _lowerCamelCase   :					float = 0.1      ,     ):



					"""simple docstring"""

					if isinstance(_lowerCamelCase      ,     _lowerCamelCase       ) and len(_lowerCamelCase       ) != batch_size:
						raise ValueError(F"""You have passed {batch_size} batch_size, but only {len(_lowerCamelCase       )} 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(_lowerCamelCase      ,     torch.Generator       ) and batch_size > 1:
						A_  :				Optional[Any]	   =		[generator] + [None] * (batch_size - 1)

					A_  :				Optional[int]	   =		[
					    ("model", self.coca_model is None),
					    ("tokenizer", self.coca_tokenizer is None),
					    ("transform", self.coca_transform is None),
					]
					A_  :				List[str]	   =		[x[0] for x in coca_is_none if x[1]]
					A_  :				Tuple	   =		", ".join(_lowerCamelCase       )
					# generate prompts with coca model if prompt is None
					if content_prompt is None:
						if len(_lowerCamelCase       ):
							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."""       )
						A_  :				Optional[Any]	   =		self.get_image_description(_lowerCamelCase       )
					if style_prompt is None:
						if len(_lowerCamelCase       ):
							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."""       )
						A_  :				Dict	   =		self.get_image_description(_lowerCamelCase       )

					# get prompt text embeddings for content and style
					A_  :				Optional[Any]	   =		self.tokenizer(
					    _lowerCamelCase      ,     padding='''max_length'''      ,     max_length=self.tokenizer.model_max_length      ,     truncation=_lowerCamelCase      ,     return_tensors='''pt'''      ,     )
					A_  :				Optional[Any]	   =		self.text_encoder(content_text_input.input_ids.to(self.device       )       )[0]

					A_  :				List[Any]	   =		self.tokenizer(
					    _lowerCamelCase      ,     padding='''max_length'''      ,     max_length=self.tokenizer.model_max_length      ,     truncation=_lowerCamelCase      ,     return_tensors='''pt'''      ,     )
					A_  :				List[str]	   =		self.text_encoder(style_text_input.input_ids.to(self.device       )       )[0]

					A_  :				Optional[Any]	   =		slerp(_lowerCamelCase      ,     _lowerCamelCase      ,     _lowerCamelCase       )

					# duplicate text embeddings for each generation per prompt
					A_  :				Optional[int]	   =		text_embeddings.repeat_interleave(_lowerCamelCase      ,     dim=0       )

					# set timesteps
					A_  :				Dict	   =		"offset" in set(inspect.signature(self.scheduler.set_timesteps       ).parameters.keys()       )
					A_  :				List[str]	   =		{}
					if accepts_offset:
						A_  :				int	   =		1

					self.scheduler.set_timesteps(_lowerCamelCase      ,     **_lowerCamelCase       )
					# 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       )

					A_  :				List[str]	   =		self.get_timesteps(_lowerCamelCase      ,     _lowerCamelCase      ,     self.device       )
					A_  :				Any	   =		timesteps[:1].repeat(_lowerCamelCase       )

					# Preprocess image
					A_  :				List[Any]	   =		preprocess(_lowerCamelCase      ,     _lowerCamelCase      ,     _lowerCamelCase       )
					A_  :				Tuple	   =		self.prepare_latents(
					    _lowerCamelCase      ,     _lowerCamelCase      ,     _lowerCamelCase      ,     text_embeddings.dtype      ,     self.device      ,     _lowerCamelCase       )

					A_  :				int	   =		preprocess(_lowerCamelCase      ,     _lowerCamelCase      ,     _lowerCamelCase       )
					A_  :				Any	   =		self.prepare_latents(
					    _lowerCamelCase      ,     _lowerCamelCase      ,     _lowerCamelCase      ,     text_embeddings.dtype      ,     self.device      ,     _lowerCamelCase       )

					A_  :				List[Any]	   =		slerp(_lowerCamelCase      ,     _lowerCamelCase      ,     _lowerCamelCase       )

					if clip_guidance_scale > 0:
						A_  :				str	   =		self.get_clip_image_embeddings(_lowerCamelCase      ,     _lowerCamelCase       )
						A_  :				Tuple	   =		self.get_clip_image_embeddings(_lowerCamelCase      ,     _lowerCamelCase       )
						A_  :				Optional[Any]	   =		slerp(
						    _lowerCamelCase      ,     _lowerCamelCase      ,     _lowerCamelCase       )

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

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

					# check if the scheduler accepts generator
					A_  :				Optional[Any]	   =		"generator" in set(inspect.signature(self.scheduler.step       ).parameters.keys()       )
					if accepts_generator:
						A_  :				Any	   =		generator

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

							# predict the noise residual
							A_  :				List[Any]	   =		self.unet(_lowerCamelCase      ,     _lowerCamelCase      ,     encoder_hidden_states=_lowerCamelCase       ).sample

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

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

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

        # Hardcode 0.18215 because stable-diffusion-2-base has not self.vae.config.scaling_factor
					A_  :				Dict	   =		1 / 0.18215 * latents
					A_  :				List[Any]	   =		self.vae.decode(_lowerCamelCase       ).sample

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

					if output_type == "pil":
						A_  :				Dict	   =		self.numpy_to_pil(_lowerCamelCase       )

					if not return_dict:
						return (image, None)

					return StableDiffusionPipelineOutput(images=_lowerCamelCase      ,     nsfw_content_detected=_lowerCamelCase       )









"""simple docstring"""

import argparse
from collections import OrderedDict
from pathlib import Path

import torch

from transformers import (
    VisualBertConfig,
    VisualBertForMultipleChoice,
    VisualBertForPreTraining,
    VisualBertForQuestionAnswering,
    VisualBertForVisualReasoning,
)
from transformers.utils import logging


logging.set_verbosity_info()
_lowerCamelCase  :			Optional[Any]        =       logging.get_logger(__name__)

_lowerCamelCase  :			Any        =       [
    ('bert.bert', 'visual_bert'),
    ('bert.cls', 'cls'),
    ('bert.classifier', 'cls'),
    ('token_type_embeddings_visual', 'visual_token_type_embeddings'),
    ('position_embeddings_visual', 'visual_position_embeddings'),
    ('projection', 'visual_projection'),
]

_lowerCamelCase  :			int        =       [
    'nlvr2_coco_pre_trained.th',
    'nlvr2_fine_tuned.th',
    'nlvr2_pre_trained.th',
    'vcr_coco_pre_train.th',
    'vcr_fine_tune.th',
    'vcr_pre_train.th',
    'vqa_coco_pre_trained.th',
    'vqa_fine_tuned.th',
    'vqa_pre_trained.th',
]

def 		lowercase_      (		_UpperCAmelCase	):




	"""simple docstring"""






	A_  :				Union[str, Any]	   =		torch.load(_UpperCAmelCase  ,      map_location='''cpu'''	)
	return sd

def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase  ,      _UpperCAmelCase=rename_keys_prefix	):




	"""simple docstring"""






	A_  :				Any	   =		OrderedDict()
	A_  :				Union[str, Any]	   =		torch.arange(config.max_position_embeddings	).expand((1, -1)	)
	# detector_d = OrderedDict()
	for key in d:
		if "detector" in key:
			# detector_d[key.replace('detector.','')] = d[key]
			continue
		A_  :				List[Any]	   =		key
		for name_pair in rename_keys_prefix:
			A_  :				Any	   =		new_key.replace(name_pair[0]  ,      name_pair[1]	)
		A_  :				Dict	   =		d[key]
		if key == "bert.cls.predictions.decoder.weight":
			# Old bert code didn't have `decoder.bias`, but was added separately
			A_  :				Any	   =		new_d['''cls.predictions.bias''']
	return new_d

@torch.no_grad()
def 		lowercase_      (		_UpperCAmelCase  ,      _UpperCAmelCase	):




	"""simple docstring"""






	assert (
	    checkpoint_path.split('''/'''	)[-1] in ACCEPTABLE_CHECKPOINTS
	), f"""The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}."""

	# Get Config
	if "pre" in checkpoint_path:
		A_  :				List[Any]	   =		'''pretraining'''
		if "vcr" in checkpoint_path:
			A_  :				Optional[Any]	   =		{'''visual_embedding_dim''': 512}
		elif "vqa_advanced" in checkpoint_path:
			A_  :				Tuple	   =		{'''visual_embedding_dim''': 2048}
		elif "vqa" in checkpoint_path:
			A_  :				str	   =		{'''visual_embedding_dim''': 2048}
		elif "nlvr" in checkpoint_path:
			A_  :				List[Any]	   =		{'''visual_embedding_dim''': 1024}
		else:
			raise NotImplementedError(f"""No implementation found for `{checkpoint_path}`."""	)
	else:
		if "vcr" in checkpoint_path:
			A_  :				Any	   =		{'''visual_embedding_dim''': 512}
			A_  :				int	   =		'''multichoice'''
		elif "vqa_advanced" in checkpoint_path:
			A_  :				Any	   =		{'''visual_embedding_dim''': 2048}
			A_  :				Any	   =		'''vqa_advanced'''
		elif "vqa" in checkpoint_path:
			A_  :				int	   =		{'''visual_embedding_dim''': 2048, '''num_labels''': 3129}
			A_  :				Any	   =		'''vqa'''
		elif "nlvr" in checkpoint_path:
			A_  :				List[Any]	   =		{
			    '''visual_embedding_dim''': 1024,
			    '''num_labels''': 2,
			}
			A_  :				List[Any]	   =		'''nlvr'''

	A_  :				Dict	   =		VisualBertConfig(**_UpperCAmelCase	)

	# Load State Dict
	A_  :				List[str]	   =		load_state_dict(_UpperCAmelCase	)

	A_  :				List[str]	   =		get_new_dict(_UpperCAmelCase  ,      _UpperCAmelCase	)

	if model_type == "pretraining":
		A_  :				List[Any]	   =		VisualBertForPreTraining(_UpperCAmelCase	)
	elif model_type == "vqa":
		A_  :				Optional[Any]	   =		VisualBertForQuestionAnswering(_UpperCAmelCase	)
	elif model_type == "nlvr":
		A_  :				Any	   =		VisualBertForVisualReasoning(_UpperCAmelCase	)
	elif model_type == "multichoice":
		A_  :				Optional[int]	   =		VisualBertForMultipleChoice(_UpperCAmelCase	)

	model.load_state_dict(_UpperCAmelCase	)
	# Save Checkpoints
	Path(_UpperCAmelCase	).mkdir(exist_ok=_UpperCAmelCase	)
	model.save_pretrained(_UpperCAmelCase	)


if __name__ == "__main__":
		_lowerCamelCase  :			Dict        =       argparse.ArgumentParser()
		# Required parameters
		parser.add_argument('orig_checkpoint_path', type=str, help='A path to .th on local filesystem.')
		parser.add_argument('pytorch_dump_folder_path', type=str, help='Path to the output PyTorch model.')
		_lowerCamelCase  :			Optional[int]        =       parser.parse_args()
		convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)





