
"""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 					__lowerCAmelCase       (_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase	):
	if isinstance(_UpperCamelCase ,						torch.Tensor	):
		return image
	elif isinstance(_UpperCamelCase ,						PIL.Image.Image	):
		__lowerCAmelCase    :		str		    =   [image]

	if isinstance(image[0] ,						PIL.Image.Image	):
		__lowerCAmelCase    :		Union[str, Any]		    =   [np.array(i.resize((w, h) ,						resample=PIL_INTERPOLATION['lanczos']	)	)[None, :] for i in image]
		__lowerCAmelCase    :		int		    =   np.concatenate(_UpperCamelCase ,						axis=0	)
		__lowerCAmelCase    :		Optional[int]		    =   np.array(_UpperCamelCase	).astype(np.floataa	) / 255.0
		__lowerCAmelCase    :		Optional[int]		    =   image.transpose(0 ,						3 ,						1 ,						2	)
		__lowerCAmelCase    :		Union[str, Any]		    =   2.0 * image - 1.0
		__lowerCAmelCase    :		int		    =   torch.from_numpy(_UpperCamelCase	)
	elif isinstance(image[0] ,						torch.Tensor	):
		__lowerCAmelCase    :		Optional[int]		    =   torch.cat(_UpperCamelCase ,						dim=0	)
	return image
def 					__lowerCAmelCase       (_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase ,						_UpperCamelCase=0.9995	):
	if not isinstance(_UpperCamelCase ,						np.ndarray	):
		__lowerCAmelCase    :		str		    =   True
		__lowerCAmelCase    :		str		    =   va.device
		__lowerCAmelCase    :		Tuple		    =   va.cpu().numpy()
		__lowerCAmelCase    :		Optional[Any]		    =   va.cpu().numpy()

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

	if inputs_are_torch:
		__lowerCAmelCase    :		Optional[Any]		    =   torch.from_numpy(_UpperCamelCase	).to(_UpperCamelCase	)

	return va
def 					__lowerCAmelCase       (_UpperCamelCase ,						_UpperCamelCase	):
	__lowerCAmelCase    :		int		    =   F.normalize(_UpperCamelCase ,						dim=-1	)
	__lowerCAmelCase    :		List[Any]		    =   F.normalize(_UpperCamelCase ,						dim=-1	)
	return (x - y).norm(dim=-1	).div(2	).arcsin().pow(2	).mul(2	)
def 					__lowerCAmelCase       (_UpperCamelCase ,						_UpperCamelCase	):
	for param in model.parameters():
		__lowerCAmelCase    :		Optional[Any]		    =   value




class 				A__ (			_lowerCamelCase):

		def __init__(   self ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE=None ,	_SCREAMING_SNAKE_CASE=None ,	_SCREAMING_SNAKE_CASE=None ,	):
			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 ,	)
			__lowerCAmelCase    :		Tuple		    =   (
			    feature_extractor.size
			    if isinstance(feature_extractor.size ,	_SCREAMING_SNAKE_CASE				)
			    else feature_extractor.size['shortest_edge']
			)
			__lowerCAmelCase    :		Optional[int]		    =   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 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE = "auto"				):
			if slice_size == "auto":
				# half the attention head size is usually a good trade-off between
				# speed and memory
				__lowerCAmelCase    :		Union[str, Any]		    =   self.unet.config.attention_head_dim // 2
			self.unet.set_attention_slice(_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			self.enable_attention_slicing(_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			set_requires_grad(self.vae ,	_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			set_requires_grad(self.vae ,	_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			set_requires_grad(self.unet ,	_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			set_requires_grad(self.unet ,	_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				):
			# get the original timestep using init_timestep
			__lowerCAmelCase    :		Union[str, Any]		    =   min(int(num_inference_steps * strength				) ,	_SCREAMING_SNAKE_CASE				)

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

			return timesteps, num_inference_steps - t_start

		def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE=None				):
			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				)}"				)

			__lowerCAmelCase    :		Any		    =   image.to(device=_SCREAMING_SNAKE_CASE ,	dtype=_SCREAMING_SNAKE_CASE				)

			if isinstance(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				):
				__lowerCAmelCase    :		str		    =   [
				    self.vae.encode(image[i : i + 1]				).latent_dist.sample(generator[i]				) for i in range(_SCREAMING_SNAKE_CASE				)
				]
				__lowerCAmelCase    :		Any		    =   torch.cat(_SCREAMING_SNAKE_CASE ,	dim=0				)
			else:
				__lowerCAmelCase    :		Dict		    =   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
			__lowerCAmelCase    :		Union[str, Any]		    =   0.1_8215 * init_latents
			__lowerCAmelCase    :		str		    =   init_latents.repeat_interleave(_SCREAMING_SNAKE_CASE ,	dim=0				)

			__lowerCAmelCase    :		Optional[int]		    =   randn_tensor(init_latents.shape ,	generator=_SCREAMING_SNAKE_CASE ,	device=_SCREAMING_SNAKE_CASE ,	dtype=_SCREAMING_SNAKE_CASE				)

			# get latents
			__lowerCAmelCase    :		List[Any]		    =   self.scheduler.add_noise(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)
			__lowerCAmelCase    :		Any		    =   init_latents

			return latents

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

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

		@torch.enable_grad()
		def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	):
			__lowerCAmelCase    :		Tuple		    =   latents.detach().requires_grad_()

			__lowerCAmelCase    :		Optional[int]		    =   self.scheduler.scale_model_input(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)

			# predict the noise residual
			__lowerCAmelCase    :		Optional[Any]		    =   self.unet(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	encoder_hidden_states=_SCREAMING_SNAKE_CASE				).sample

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

				__lowerCAmelCase    :		str		    =   torch.sqrt(_SCREAMING_SNAKE_CASE				)
				__lowerCAmelCase    :		List[str]		    =   pred_original_sample * (fac) + latents * (1 - fac)
			elif isinstance(self.scheduler ,	_SCREAMING_SNAKE_CASE				):
				__lowerCAmelCase    :		Union[str, Any]		    =   self.scheduler.sigmas[index]
				__lowerCAmelCase    :		List[str]		    =   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
			__lowerCAmelCase    :		str		    =   1 / 0.1_8215 * sample
			__lowerCAmelCase    :		int		    =   self.vae.decode(_SCREAMING_SNAKE_CASE				).sample
			__lowerCAmelCase    :		Tuple		    =   (image / 2 + 0.5).clamp(0 ,	1				)

			__lowerCAmelCase    :		Any		    =   transforms.Resize(self.feature_extractor_size				)(_SCREAMING_SNAKE_CASE				)
			__lowerCAmelCase    :		Tuple		    =   self.normalize(_SCREAMING_SNAKE_CASE				).to(latents.dtype				)

			__lowerCAmelCase    :		Optional[int]		    =   self.clip_model.get_image_features(_SCREAMING_SNAKE_CASE				)
			__lowerCAmelCase    :		Any		    =   image_embeddings_clip / image_embeddings_clip.norm(p=2 ,	dim=-1 ,	keepdim=_SCREAMING_SNAKE_CASE				)

			__lowerCAmelCase    :		int		    =   spherical_dist_loss(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				).mean() * clip_guidance_scale

			__lowerCAmelCase    :		List[str]		    =   -torch.autograd.grad(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)[0]

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







		@torch.no_grad()
		def __call__(   self ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE = None ,	_SCREAMING_SNAKE_CASE = None ,	_SCREAMING_SNAKE_CASE = 5_12 ,	_SCREAMING_SNAKE_CASE = 5_12 ,	_SCREAMING_SNAKE_CASE = 0.6 ,	_SCREAMING_SNAKE_CASE = 50 ,	_SCREAMING_SNAKE_CASE = 7.5 ,	_SCREAMING_SNAKE_CASE = 1 ,	_SCREAMING_SNAKE_CASE = 0.0 ,	_SCREAMING_SNAKE_CASE = 1_00 ,	_SCREAMING_SNAKE_CASE = None ,	_SCREAMING_SNAKE_CASE = "pil" ,	_SCREAMING_SNAKE_CASE = True ,	_SCREAMING_SNAKE_CASE = 0.8 ,	_SCREAMING_SNAKE_CASE = 0.1 ,	_SCREAMING_SNAKE_CASE = 0.1 ,	):
			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:
				__lowerCAmelCase    :		Optional[Any]		    =   [generator] + [None] * (batch_size - 1)

			__lowerCAmelCase    :		Optional[int]		    =   [
			    ('model', self.coca_model is None),
			    ('tokenizer', self.coca_tokenizer is None),
			    ('transform', self.coca_transform is None),
			]
			__lowerCAmelCase    :		Union[str, Any]		    =   [x[0] for x in coca_is_none if x[1]]
			__lowerCAmelCase    :		Tuple		    =   ', '.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."				)
				__lowerCAmelCase    :		List[Any]		    =   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."				)
				__lowerCAmelCase    :		Optional[Any]		    =   self.get_image_description(_SCREAMING_SNAKE_CASE				)

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

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

			__lowerCAmelCase    :		str		    =   slerp(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)

			# duplicate text embeddings for each generation per prompt
			__lowerCAmelCase    :		List[str]		    =   text_embeddings.repeat_interleave(_SCREAMING_SNAKE_CASE ,	dim=0				)

			# set timesteps
			__lowerCAmelCase    :		Union[str, Any]		    =   'offset' in set(inspect.signature(self.scheduler.set_timesteps				).parameters.keys()				)
			__lowerCAmelCase    :		List[Any]		    =   {}
			if accepts_offset:
				__lowerCAmelCase    :		Union[str, Any]		    =   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				)

			__lowerCAmelCase    ,     __lowerCAmelCase    :		int		    =   self.get_timesteps(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	self.device				)
			__lowerCAmelCase    :		Any		    =   timesteps[:1].repeat(_SCREAMING_SNAKE_CASE				)

			# Preprocess image
			__lowerCAmelCase    :		Union[str, Any]		    =   preprocess(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)
			__lowerCAmelCase    :		List[Any]		    =   self.prepare_latents(
			    _SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	text_embeddings.dtype ,	self.device ,	_SCREAMING_SNAKE_CASE				)

			__lowerCAmelCase    :		Tuple		    =   preprocess(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)
			__lowerCAmelCase    :		str		    =   self.prepare_latents(
			    _SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	text_embeddings.dtype ,	self.device ,	_SCREAMING_SNAKE_CASE				)

			__lowerCAmelCase    :		Tuple		    =   slerp(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)

			if clip_guidance_scale > 0:
				__lowerCAmelCase    :		Any		    =   self.get_clip_image_embeddings(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)
				__lowerCAmelCase    :		Union[str, Any]		    =   self.get_clip_image_embeddings(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)
				__lowerCAmelCase    :		Any		    =   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.
			__lowerCAmelCase    :		int		    =   guidance_scale > 1.0
			# get unconditional embeddings for classifier free guidance
			if do_classifier_free_guidance:
				__lowerCAmelCase    :		int		    =   content_text_input.input_ids.shape[-1]
				__lowerCAmelCase    :		Optional[Any]		    =   self.tokenizer([''] ,	padding='max_length' ,	max_length=_SCREAMING_SNAKE_CASE ,	return_tensors='pt'				)
				__lowerCAmelCase    :		Tuple		    =   self.text_encoder(uncond_input.input_ids.to(self.device				)				)[0]
				# duplicate unconditional embeddings for each generation per prompt
				__lowerCAmelCase    :		int		    =   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
				__lowerCAmelCase    :		Tuple		    =   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    :		Optional[int]		    =   (batch_size, self.unet.config.in_channels, height // 8, width // 8)
			__lowerCAmelCase    :		int		    =   text_embeddings.dtype
			if latents is None:
				if self.device.type == "mps":
					# randn does not work reproducibly on mps
					__lowerCAmelCase    :		Optional[Any]		    =   torch.randn(_SCREAMING_SNAKE_CASE ,	generator=_SCREAMING_SNAKE_CASE ,	device='cpu' ,	dtype=_SCREAMING_SNAKE_CASE				).to(
					    self.device				)
				else:
					__lowerCAmelCase    :		int		    =   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}"				)
				__lowerCAmelCase    :		Any		    =   latents.to(self.device				)

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

			# check if the scheduler accepts generator
			__lowerCAmelCase    :		List[Any]		    =   'generator' in set(inspect.signature(self.scheduler.step				).parameters.keys()				)
			if accepts_generator:
				__lowerCAmelCase    :		Any		    =   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
					__lowerCAmelCase    :		Any		    =   torch.cat([latents] * 2				) if do_classifier_free_guidance else latents
					__lowerCAmelCase    :		Optional[int]		    =   self.scheduler.scale_model_input(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)

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

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

					# perform clip guidance
					if clip_guidance_scale > 0:
						__lowerCAmelCase    :		Optional[int]		    =   (
						    text_embeddings.chunk(2				)[1] if do_classifier_free_guidance else text_embeddings
						)
						__lowerCAmelCase    ,     __lowerCAmelCase    :		Any		    =   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
					__lowerCAmelCase    :		Tuple		    =   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
			__lowerCAmelCase    :		int		    =   1 / 0.1_8215 * latents
			__lowerCAmelCase    :		Union[str, Any]		    =   self.vae.decode(_SCREAMING_SNAKE_CASE				).sample

			__lowerCAmelCase    :		Dict		    =   (image / 2 + 0.5).clamp(0 ,	1				)
			__lowerCAmelCase    :		int		    =   image.cpu().permute(0 ,	2 ,	3 ,	1				).numpy()

			if output_type == "pil":
				__lowerCAmelCase    :		Optional[Any]		    =   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 RegNetConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

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 RegNetForImageClassification, RegNetModel
							from transformers.models.regnet.modeling_regnet import REGNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor




class 				A__ :

		def __init__(   self ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE=3 ,	_SCREAMING_SNAKE_CASE=32 ,	_SCREAMING_SNAKE_CASE=3 ,	_SCREAMING_SNAKE_CASE=10 ,	_SCREAMING_SNAKE_CASE=[10, 20, 30, 40] ,	_SCREAMING_SNAKE_CASE=[1, 1, 2, 1] ,	_SCREAMING_SNAKE_CASE=True ,	_SCREAMING_SNAKE_CASE=True ,	_SCREAMING_SNAKE_CASE="relu" ,	_SCREAMING_SNAKE_CASE=3 ,	_SCREAMING_SNAKE_CASE=None ,	):
			__lowerCAmelCase    :		Optional[Any]		    =   parent
			__lowerCAmelCase    :		int		    =   batch_size
			__lowerCAmelCase    :		Union[str, Any]		    =   image_size
			__lowerCAmelCase    :		Optional[Any]		    =   num_channels
			__lowerCAmelCase    :		List[Any]		    =   embeddings_size
			__lowerCAmelCase    :		Dict		    =   hidden_sizes
			__lowerCAmelCase    :		Optional[int]		    =   depths
			__lowerCAmelCase    :		Optional[Any]		    =   is_training
			__lowerCAmelCase    :		int		    =   use_labels
			__lowerCAmelCase    :		List[Any]		    =   hidden_act
			__lowerCAmelCase    :		Optional[Any]		    =   num_labels
			__lowerCAmelCase    :		List[str]		    =   scope
			__lowerCAmelCase    :		Optional[Any]		    =   len(_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		Optional[Any]		    =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

			__lowerCAmelCase    :		List[str]		    =   None
			if self.use_labels:
				__lowerCAmelCase    :		Tuple		    =   ids_tensor([self.batch_size] ,	self.num_labels				)

			__lowerCAmelCase    :		int		    =   self.get_config()

			return config, pixel_values, labels

		def 							__lowerCamelCase							(   self				):
			return RegNetConfig(
			    num_channels=self.num_channels ,	embeddings_size=self.embeddings_size ,	hidden_sizes=self.hidden_sizes ,	depths=self.depths ,	hidden_act=self.hidden_act ,	num_labels=self.num_labels ,	)

		def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				):
			__lowerCAmelCase    :		Any		    =   RegNetModel(config=_SCREAMING_SNAKE_CASE				)
			model.to(_SCREAMING_SNAKE_CASE				)
			model.eval()
			__lowerCAmelCase    :		Any		    =   model(_SCREAMING_SNAKE_CASE				)
			# expected last hidden states: B, C, H // 32, W // 32
			self.parent.assertEqual(
			    result.last_hidden_state.shape ,	(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32) ,	)

		def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				):
			__lowerCAmelCase    :		Optional[Any]		    =   self.num_labels
			__lowerCAmelCase    :		str		    =   RegNetForImageClassification(_SCREAMING_SNAKE_CASE				)
			model.to(_SCREAMING_SNAKE_CASE				)
			model.eval()
			__lowerCAmelCase    :		Any		    =   model(_SCREAMING_SNAKE_CASE ,	labels=_SCREAMING_SNAKE_CASE				)
			self.parent.assertEqual(result.logits.shape ,	(self.batch_size, self.num_labels)				)







		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		Optional[Any]		    =   self.prepare_config_and_inputs()
			__lowerCAmelCase    ,     __lowerCAmelCase    ,     __lowerCAmelCase    :		Any		    =   config_and_inputs
			__lowerCAmelCase    :		List[str]		    =   {'pixel_values': pixel_values}
			return config, inputs_dict




@require_torch
class 				A__ (			_lowerCamelCase ,  _lowerCamelCase ,  unittest.TestCase):
		A_		: List[str]            =			(RegNetModel, RegNetForImageClassification) if is_torch_available() else ()
		A_		: List[Any]            =			(
		    {'feature-extraction': RegNetModel, 'image-classification': RegNetForImageClassification}
		    if is_torch_available()
		    else {}
		)

		A_		: Dict            =			False
		A_		: str            =			False
		A_		: int            =			False
		A_		: Union[str, Any]            =			False

		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		List[str]		    =   RegNetModelTester(self				)
			__lowerCAmelCase    :		int		    =   ConfigTester(self ,	config_class=_SCREAMING_SNAKE_CASE ,	has_text_modality=_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			self.create_and_test_config_common_properties()
			self.config_tester.create_and_test_config_to_json_string()
			self.config_tester.create_and_test_config_to_json_file()
			self.config_tester.create_and_test_config_from_and_save_pretrained()
			self.config_tester.create_and_test_config_with_num_labels()
			self.config_tester.check_config_can_be_init_without_params()
			self.config_tester.check_config_arguments_init()

		def 							__lowerCamelCase							(   self				):
			return

		@unittest.skip(reason='RegNet does not use inputs_embeds'				)
		def 							__lowerCamelCase							(   self				):
			pass

		@unittest.skip(reason='RegNet does not support input and output embeddings'				)
		def 							__lowerCamelCase							(   self				):
			pass

		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    ,     __lowerCAmelCase    :		List[Any]		    =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
				__lowerCAmelCase    :		Optional[int]		    =   model_class(_SCREAMING_SNAKE_CASE				)
				__lowerCAmelCase    :		Dict		    =   inspect.signature(model.forward				)
				# signature.parameters is an OrderedDict => so arg_names order is deterministic
				__lowerCAmelCase    :		Tuple		    =   [*signature.parameters.keys()]

				__lowerCAmelCase    :		Optional[int]		    =   ['pixel_values']
				self.assertListEqual(arg_names[:1] ,	_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		Optional[Any]		    =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    ,     __lowerCAmelCase    :		List[str]		    =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
				__lowerCAmelCase    :		Union[str, Any]		    =   model_class(config=_SCREAMING_SNAKE_CASE				)
				for name, module in model.named_modules():
					if isinstance(_SCREAMING_SNAKE_CASE ,	(nn.BatchNormad, nn.GroupNorm)				):
						self.assertTrue(
						    torch.all(module.weight == 1				) ,	msg=f"Parameter {name} of model {model_class} seems not properly initialized" ,	)
						self.assertTrue(
						    torch.all(module.bias == 0				) ,	msg=f"Parameter {name} of model {model_class} seems not properly initialized" ,	)

		def 							__lowerCamelCase							(   self				):
			def check_hidden_states_output(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				):
				__lowerCAmelCase    :		Any		    =   model_class(_SCREAMING_SNAKE_CASE				)
				model.to(_SCREAMING_SNAKE_CASE				)
				model.eval()

				with torch.no_grad():
					__lowerCAmelCase    :		List[str]		    =   model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)				)

				__lowerCAmelCase    :		Any		    =   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

				__lowerCAmelCase    :		List[str]		    =   self.model_tester.num_stages
				self.assertEqual(len(_SCREAMING_SNAKE_CASE				) ,	expected_num_stages + 1				)

				# RegNet's feature maps are of shape (batch_size, num_channels, height, width)
				self.assertListEqual(
				    list(hidden_states[0].shape[-2:]				) ,	[self.model_tester.image_size // 2, self.model_tester.image_size // 2] ,	)

			__lowerCAmelCase    ,     __lowerCAmelCase    :		Dict		    =   self.model_tester.prepare_config_and_inputs_for_common()
			__lowerCAmelCase    :		List[Any]		    =   ['basic', 'bottleneck']
			for model_class in self.all_model_classes:
				for layer_type in layers_type:
					__lowerCAmelCase    :		List[Any]		    =   layer_type
					__lowerCAmelCase    :		Tuple		    =   True
					check_hidden_states_output(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					__lowerCAmelCase    :		Dict		    =   True

					check_hidden_states_output(_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE				)

		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		List[Any]		    =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE				)







		@slow
		def 							__lowerCamelCase							(   self				):
			for model_name in REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
				__lowerCAmelCase    :		Dict		    =   RegNetModel.from_pretrained(_SCREAMING_SNAKE_CASE				)
				self.assertIsNotNone(_SCREAMING_SNAKE_CASE				)



def 					__lowerCAmelCase       ():
	__lowerCAmelCase    :		str		    =   Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
	return image




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

		@cached_property
		def 							__lowerCamelCase							(   self				):
			return (
			    AutoImageProcessor.from_pretrained(REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]				)
			    if is_vision_available()
			    else None
			)







		@slow
		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		List[str]		    =   RegNetForImageClassification.from_pretrained(REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]				).to(_SCREAMING_SNAKE_CASE				)

			__lowerCAmelCase    :		str		    =   self.default_image_processor
			__lowerCAmelCase    :		Any		    =   prepare_img()
			__lowerCAmelCase    :		Any		    =   image_processor(images=_SCREAMING_SNAKE_CASE ,	return_tensors='pt'				).to(_SCREAMING_SNAKE_CASE				)

			# forward pass
			with torch.no_grad():
				__lowerCAmelCase    :		List[str]		    =   model(**_SCREAMING_SNAKE_CASE				)

			# verify the logits
			__lowerCAmelCase    :		Union[str, Any]		    =   torch.Size((1, 10_00)				)
			self.assertEqual(outputs.logits.shape ,	_SCREAMING_SNAKE_CASE				)

			__lowerCAmelCase    :		Optional[int]		    =   torch.tensor([-0.4180, -1.5051, -3.4836]				).to(_SCREAMING_SNAKE_CASE				)

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