


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 	lowerCamelCase_							(       UpperCAmelCase__							,     UpperCAmelCase__							,     UpperCAmelCase__   ):



						"""simple docstring"""




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

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




def 	lowerCamelCase_							(       UpperCAmelCase__							,     UpperCAmelCase__							,     UpperCAmelCase__							,     UpperCAmelCase__=0.99_95   ):



						"""simple docstring"""




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

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

						if inputs_are_torch:
												a_		      =		torch.from_numpy(_lowerCamelCase   ).to(_lowerCamelCase   )

						return va




def 	lowerCamelCase_							(       UpperCAmelCase__							,     UpperCAmelCase__   ):



						"""simple docstring"""




						a_		      =		F.normalize(_lowerCamelCase							,     dim=-1   )
						a_		      =		F.normalize(_lowerCamelCase							,     dim=-1   )
						return (x - y).norm(dim=-1   ).div(2   ).arcsin().pow(2   ).mul(2   )




def 	lowerCamelCase_							(       UpperCAmelCase__							,     UpperCAmelCase__   ):



						"""simple docstring"""




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





class 							lowercase_					(			__snake_case):






							"""simple docstring"""







							def __init__( self      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase=None      ,     _UpperCAmelCase=None      ,     _UpperCAmelCase=None      ,     ):







													"""simple docstring"""


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


							def     lowercase__  ( self      ,     _UpperCAmelCase = "auto"		):







													"""simple docstring"""


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


							def     lowercase__  ( self		):







													"""simple docstring"""


													self.enable_attention_slicing(__SCREAMING_SNAKE_CASE		)


							def     lowercase__  ( self		):







													"""simple docstring"""


													set_requires_grad(self.vae      ,     __SCREAMING_SNAKE_CASE		)


							def     lowercase__  ( self		):







													"""simple docstring"""


													set_requires_grad(self.vae      ,     __SCREAMING_SNAKE_CASE		)


							def     lowercase__  ( self		):







													"""simple docstring"""


													set_requires_grad(self.unet      ,     __SCREAMING_SNAKE_CASE		)


							def     lowercase__  ( self		):







													"""simple docstring"""


													set_requires_grad(self.unet      ,     __SCREAMING_SNAKE_CASE		)


							def     lowercase__  ( self      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase		):







													"""simple docstring"""


													a_		      =		min(int(num_inference_steps * strength		)      ,     __SCREAMING_SNAKE_CASE		)

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

													return timesteps, num_inference_steps - t_start


							def     lowercase__  ( self      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase=None		):







													"""simple docstring"""


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

													a_		      =		image.to(device=__SCREAMING_SNAKE_CASE      ,     dtype=__SCREAMING_SNAKE_CASE		)

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

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

													a_		      =		randn_tensor(init_latents.shape      ,     generator=__SCREAMING_SNAKE_CASE      ,     device=__SCREAMING_SNAKE_CASE      ,     dtype=__SCREAMING_SNAKE_CASE		)

													# get latents
													a_		      =		self.scheduler.add_noise(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE		)
													a_		      =		init_latents

													return latents


							def     lowercase__  ( self      ,     _UpperCAmelCase		):







													"""simple docstring"""


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


							def     lowercase__  ( self      ,     _UpperCAmelCase      ,     _UpperCAmelCase		):







													"""simple docstring"""


													a_		      =		self.feature_extractor.preprocess(__SCREAMING_SNAKE_CASE		)
													a_		      =		torch.from_numpy(clip_image_input["""pixel_values"""][0]		).unsqueeze(0		).to(self.device		).half()
													a_		      =		self.clip_model.get_image_features(__SCREAMING_SNAKE_CASE		)
													a_		      =		image_embeddings_clip / image_embeddings_clip.norm(p=2      ,     dim=-1      ,     keepdim=__SCREAMING_SNAKE_CASE		)
													a_		      =		image_embeddings_clip.repeat_interleave(__SCREAMING_SNAKE_CASE      ,     dim=0		)
													return image_embeddings_clip


							@torch.enable_grad()
							def     lowercase__  ( self      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     ):







													"""simple docstring"""


													a_		      =		latents.detach().requires_grad_()

													a_		      =		self.scheduler.scale_model_input(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE		)

													# predict the noise residual
													a_		      =		self.unet(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     encoder_hidden_states=__SCREAMING_SNAKE_CASE		).sample

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

																			a_		      =		torch.sqrt(__SCREAMING_SNAKE_CASE		)
																			a_		      =		pred_original_sample * (fac) + latents * (1 - fac)
													elif isinstance(self.scheduler      ,     __SCREAMING_SNAKE_CASE		):
																			a_		      =		self.scheduler.sigmas[index]
																			a_		      =		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_		      =		1 / 0.1_8_2_1_5 * sample
													a_		      =		self.vae.decode(__SCREAMING_SNAKE_CASE		).sample
													a_		      =		(image / 2 + 0.5).clamp(0      ,     1		)

													a_		      =		transforms.Resize(self.feature_extractor_size		)(__SCREAMING_SNAKE_CASE		)
													a_		      =		self.normalize(__SCREAMING_SNAKE_CASE		).to(latents.dtype		)

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

													a_		      =		spherical_dist_loss(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE		).mean() * clip_guidance_scale

													a_		      =		-torch.autograd.grad(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE		)[0]

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






							@torch.no_grad()
							def __call__( self      ,     _UpperCAmelCase      ,     _UpperCAmelCase      ,     _UpperCAmelCase = None      ,     _UpperCAmelCase = None      ,     _UpperCAmelCase = 512      ,     _UpperCAmelCase = 512      ,     _UpperCAmelCase = 0.6      ,     _UpperCAmelCase = 50      ,     _UpperCAmelCase = 7.5      ,     _UpperCAmelCase = 1      ,     _UpperCAmelCase = 0.0      ,     _UpperCAmelCase = 100      ,     _UpperCAmelCase = None      ,     _UpperCAmelCase = "pil"      ,     _UpperCAmelCase = True      ,     _UpperCAmelCase = 0.8      ,     _UpperCAmelCase = 0.1      ,     _UpperCAmelCase = 0.1      ,     ):







													"""simple docstring"""


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

													a_		      =		[
													    ("""model""", self.coca_model is None),
													    ("""tokenizer""", self.coca_tokenizer is None),
													    ("""transform""", self.coca_transform is None),
													]
													a_		      =		[x[0] for x in coca_is_none if x[1]]
													a_		      =		""", """.join(__SCREAMING_SNAKE_CASE		)
													# generate prompts with coca model if prompt is None
													if content_prompt is None:
																			if len(__SCREAMING_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."		)
																			a_		      =		self.get_image_description(__SCREAMING_SNAKE_CASE		)
													if style_prompt is None:
																			if len(__SCREAMING_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."		)
																			a_		      =		self.get_image_description(__SCREAMING_SNAKE_CASE		)

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

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

													a_		      =		slerp(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE		)

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

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

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

													a_      , a_		      =		self.get_timesteps(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     self.device		)
													a_		      =		timesteps[:1].repeat(__SCREAMING_SNAKE_CASE		)

													# Preprocess image
													a_		      =		preprocess(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE		)
													a_		      =		self.prepare_latents(
													    __SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     text_embeddings.dtype      ,     self.device      ,     __SCREAMING_SNAKE_CASE		)

													a_		      =		preprocess(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE		)
													a_		      =		self.prepare_latents(
													    __SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     text_embeddings.dtype      ,     self.device      ,     __SCREAMING_SNAKE_CASE		)

													a_		      =		slerp(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE		)

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

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

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

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

																									# predict the noise residual
																									a_		      =		self.unet(__SCREAMING_SNAKE_CASE      ,     __SCREAMING_SNAKE_CASE      ,     encoder_hidden_states=__SCREAMING_SNAKE_CASE		).sample

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

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

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

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

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

													if output_type == "pil":
																			a_		      =		self.numpy_to_pil(__SCREAMING_SNAKE_CASE		)

													if not return_dict:
																			return (image, None)

													return StableDiffusionPipelineOutput(images=__SCREAMING_SNAKE_CASE      ,     nsfw_content_detected=__SCREAMING_SNAKE_CASE		)







"""simple docstring"""





import inspect
import unittest

from transformers import ViTMSNConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import ViTMSNForImageClassification, ViTMSNModel
  from transformers.models.vit_msn.modeling_vit_msn import VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import ViTImageProcessor






class 	a		:






    def __init__(			self	:				Union[str, Any]			,		__SCREAMING_SNAKE_CASE	:				Tuple			,		__SCREAMING_SNAKE_CASE	:				str=13			,		__SCREAMING_SNAKE_CASE	:				int=30			,		__SCREAMING_SNAKE_CASE	:				List[Any]=2			,		__SCREAMING_SNAKE_CASE	:				Dict=3			,		__SCREAMING_SNAKE_CASE	:				List[str]=True			,		__SCREAMING_SNAKE_CASE	:				str=True			,		__SCREAMING_SNAKE_CASE	:				Any=32			,		__SCREAMING_SNAKE_CASE	:				List[str]=5			,		__SCREAMING_SNAKE_CASE	:				List[str]=4			,		__SCREAMING_SNAKE_CASE	:				Union[str, Any]=37			,		__SCREAMING_SNAKE_CASE	:				Union[str, Any]="gelu"			,		__SCREAMING_SNAKE_CASE	:				List[Any]=0.1			,		__SCREAMING_SNAKE_CASE	:				Tuple=0.1			,		__SCREAMING_SNAKE_CASE	:				Optional[Any]=10			,		__SCREAMING_SNAKE_CASE	:				Optional[Any]=0.02			,		__SCREAMING_SNAKE_CASE	:				Union[str, Any]=None			,		)	->						List[str]:
      lowerCamelCase_       	=		parent
      lowerCamelCase_       	=		batch_size
      lowerCamelCase_       	=		image_size
      lowerCamelCase_       	=		patch_size
      lowerCamelCase_       	=		num_channels
      lowerCamelCase_       	=		is_training
      lowerCamelCase_       	=		use_labels
      lowerCamelCase_       	=		hidden_size
      lowerCamelCase_       	=		num_hidden_layers
      lowerCamelCase_       	=		num_attention_heads
      lowerCamelCase_       	=		intermediate_size
      lowerCamelCase_       	=		hidden_act
      lowerCamelCase_       	=		hidden_dropout_prob
      lowerCamelCase_       	=		attention_probs_dropout_prob
      lowerCamelCase_       	=		type_sequence_label_size
      lowerCamelCase_       	=		initializer_range
      lowerCamelCase_       	=		scope

      # in ViT MSN, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
      lowerCamelCase_       	=		(image_size // patch_size) ** 2
      lowerCamelCase_       	=		num_patches + 1






    def   UpperCamelCase			(			self	:				Dict   )	->						Union[str, Any]:
      lowerCamelCase_       	=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

      lowerCamelCase_       	=		None
      if self.use_labels:
        lowerCamelCase_       	=		ids_tensor([self.batch_size]			,		self.type_sequence_label_size   )

      lowerCamelCase_       	=		self.get_config()

      return config, pixel_values, labels






    def   UpperCamelCase			(			self	:				Union[str, Any]   )	->						Union[str, Any]:
      return ViTMSNConfig(
          image_size=self.image_size			,		patch_size=self.patch_size			,		num_channels=self.num_channels			,		hidden_size=self.hidden_size			,		num_hidden_layers=self.num_hidden_layers			,		num_attention_heads=self.num_attention_heads			,		intermediate_size=self.intermediate_size			,		hidden_act=self.hidden_act			,		hidden_dropout_prob=self.hidden_dropout_prob			,		attention_probs_dropout_prob=self.attention_probs_dropout_prob			,		initializer_range=self.initializer_range			,		)






    def   UpperCamelCase			(			self	:				int			,		__SCREAMING_SNAKE_CASE	:				Any			,		__SCREAMING_SNAKE_CASE	:				Any			,		__SCREAMING_SNAKE_CASE	:				int   )	->						Any:
      lowerCamelCase_       	=		ViTMSNModel(config=__SCREAMING_SNAKE_CASE   )
      model.to(__SCREAMING_SNAKE_CASE   )
      model.eval()
      lowerCamelCase_       	=		model(__SCREAMING_SNAKE_CASE   )
      self.parent.assertEqual(result.last_hidden_state.shape			,		(self.batch_size, self.seq_length, self.hidden_size)   )






    def   UpperCamelCase			(			self	:				int			,		__SCREAMING_SNAKE_CASE	:				Dict			,		__SCREAMING_SNAKE_CASE	:				str			,		__SCREAMING_SNAKE_CASE	:				str   )	->						List[str]:
      lowerCamelCase_       	=		self.type_sequence_label_size
      lowerCamelCase_       	=		ViTMSNForImageClassification(__SCREAMING_SNAKE_CASE   )
      model.to(__SCREAMING_SNAKE_CASE   )
      model.eval()
      lowerCamelCase_       	=		model(__SCREAMING_SNAKE_CASE			,		labels=__SCREAMING_SNAKE_CASE   )
      print('Pixel and labels shape: {pixel_values.shape}, {labels.shape}'   )
      print('Labels: {labels}'   )
      self.parent.assertEqual(result.logits.shape			,		(self.batch_size, self.type_sequence_label_size)   )

      # test greyscale images
      lowerCamelCase_       	=		1
      lowerCamelCase_       	=		ViTMSNForImageClassification(__SCREAMING_SNAKE_CASE   )
      model.to(__SCREAMING_SNAKE_CASE   )
      model.eval()

      lowerCamelCase_       	=		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
      lowerCamelCase_       	=		model(__SCREAMING_SNAKE_CASE   )
      self.parent.assertEqual(result.logits.shape			,		(self.batch_size, self.type_sequence_label_size)   )






    def   UpperCamelCase			(			self	:				List[str]   )	->						Tuple:
      lowerCamelCase_       	=		self.prepare_config_and_inputs()
      lowerCamelCase_	,      lowerCamelCase_	,      lowerCamelCase_       	=		config_and_inputs
      lowerCamelCase_       	=		{'pixel_values': pixel_values}
      return config, inputs_dict






@require_torch
class 	a		(   __snake_case    , __snake_case    , unittest.TestCase						):
    SCREAMING_SNAKE_CASE		:	str		=					(ViTMSNModel, ViTMSNForImageClassification) if is_torch_available() else ()
    SCREAMING_SNAKE_CASE		:	Dict		=					(
        {"""feature-extraction""": ViTMSNModel, """image-classification""": ViTMSNForImageClassification}
        if is_torch_available()
        else {}
    )

    SCREAMING_SNAKE_CASE		:	Dict		=					False
    SCREAMING_SNAKE_CASE		:	Union[str, Any]		=					False
    SCREAMING_SNAKE_CASE		:	Optional[Any]		=					False
    SCREAMING_SNAKE_CASE		:	str		=					False






    def   UpperCamelCase			(			self	:				List[Any]   )	->						List[str]:
      lowerCamelCase_       	=		ViTMSNModelTester(self   )
      lowerCamelCase_       	=		ConfigTester(self			,		config_class=__SCREAMING_SNAKE_CASE			,		has_text_modality=__SCREAMING_SNAKE_CASE			,		hidden_size=37   )






    def   UpperCamelCase			(			self	:				Union[str, Any]   )	->						Optional[Any]:
      self.config_tester.run_common_tests()






    @unittest.skip(reason='ViTMSN does not use inputs_embeds'   )
    def   UpperCamelCase			(			self	:				Tuple   )	->						Union[str, Any]:
      pass






    def   UpperCamelCase			(			self	:				List[Any]   )	->						List[Any]:
      lowerCamelCase_	,      lowerCamelCase_       	=		self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        lowerCamelCase_       	=		model_class(__SCREAMING_SNAKE_CASE   )
        self.assertIsInstance(model.get_input_embeddings()			,		(nn.Module)   )
        lowerCamelCase_       	=		model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(__SCREAMING_SNAKE_CASE			,		nn.Linear   )   )






    def   UpperCamelCase			(			self	:				Optional[int]   )	->						Any:
      lowerCamelCase_	,      lowerCamelCase_       	=		self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        lowerCamelCase_       	=		model_class(__SCREAMING_SNAKE_CASE   )
        lowerCamelCase_       	=		inspect.signature(model.forward   )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        lowerCamelCase_       	=		[*signature.parameters.keys()]

        lowerCamelCase_       	=		['pixel_values']
        self.assertListEqual(arg_names[:1]			,		__SCREAMING_SNAKE_CASE   )






    def   UpperCamelCase			(			self	:				Optional[Any]   )	->						List[Any]:
      lowerCamelCase_       	=		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE   )






    def   UpperCamelCase			(			self	:				Union[str, Any]   )	->						Optional[int]:
      lowerCamelCase_       	=		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*__SCREAMING_SNAKE_CASE   )






    @slow
    def   UpperCamelCase			(			self	:				int   )	->						List[Any]:
      for model_name in VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowerCamelCase_       	=		ViTMSNModel.from_pretrained(__SCREAMING_SNAKE_CASE   )
        self.assertIsNotNone(__SCREAMING_SNAKE_CASE   )







def 	lowerCamelCase__			(				) ->							List[Any]:
  lowerCamelCase_       	=		Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'					)
  return image






@require_torch
@require_vision
class 	a		(   unittest.TestCase						):






    @cached_property
    def   UpperCamelCase			(			self	:				Optional[int]   )	->						Any:
      return ViTImageProcessor.from_pretrained('facebook/vit-msn-small'   ) if is_vision_available() else None






    @slow
    def   UpperCamelCase			(			self	:				Dict   )	->						Any:
      torch.manual_seed(2   )
      lowerCamelCase_       	=		ViTMSNForImageClassification.from_pretrained('facebook/vit-msn-small'   ).to(__SCREAMING_SNAKE_CASE   )

      lowerCamelCase_       	=		self.default_image_processor
      lowerCamelCase_       	=		prepare_img()
      lowerCamelCase_       	=		image_processor(images=__SCREAMING_SNAKE_CASE			,		return_tensors='pt'   ).to(__SCREAMING_SNAKE_CASE   )

      # forward pass
      with torch.no_grad():
        lowerCamelCase_       	=		model(**__SCREAMING_SNAKE_CASE   )

      # verify the logits
      lowerCamelCase_       	=		torch.Size((1, 1000)   )
      self.assertEqual(outputs.logits.shape			,		__SCREAMING_SNAKE_CASE   )

      lowerCamelCase_       	=		torch.tensor([-0.0_803, -0.4_454, -0.2_375]   ).to(__SCREAMING_SNAKE_CASE   )

      self.assertTrue(torch.allclose(outputs.logits[0, :3]			,		__SCREAMING_SNAKE_CASE			,		atol=1e-4   )   )




