






'''simple docstring'''







import inspect
from typing import Callable, List, Optional, Union

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


UpperCamelCase    :      Optional[int]  			=		logging.get_logger(__name__)  # pylint: disable=invalid-name




class 		UpperCamelCase__    (a			):




		'''simple docstring'''





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






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






		def 			UpperCamelCase_					(      self       ):
								self.enable_attention_slicing(_lowerCAmelCase       )






		@torch.no_grad()
		def __call__(      self   ,_lowerCAmelCase   ,_lowerCAmelCase = 5_12   ,_lowerCAmelCase = 5_12   ,_lowerCAmelCase = 50   ,_lowerCAmelCase = 7.5   ,_lowerCAmelCase = None   ,_lowerCAmelCase = 1   ,_lowerCAmelCase = 0.0   ,_lowerCAmelCase = None   ,_lowerCAmelCase = None   ,_lowerCAmelCase = "pil"   ,_lowerCAmelCase = True   ,_lowerCAmelCase = None   ,_lowerCAmelCase = 1   ,_lowerCAmelCase = None   ,**_lowerCAmelCase   ,):

								if isinstance(_lowerCAmelCase   ,_lowerCAmelCase       ):
														lowerCamelCase__         =  1
								elif isinstance(_lowerCAmelCase   ,_lowerCAmelCase       ):
														lowerCamelCase__         =  len(_lowerCAmelCase       )
								else:
														raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(_lowerCAmelCase       )}'''       )

								if height % 8 != 0 or width % 8 != 0:
														raise ValueError(F'''`height` and `width` have to be divisible by 8 but are {height} and {width}.'''       )

								if (callback_steps is None) or (
								    callback_steps is not None and (not isinstance(_lowerCAmelCase   ,_lowerCAmelCase       ) or callback_steps <= 0)
								):
														raise ValueError(
														    F'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
														    F''' {type(_lowerCAmelCase       )}.'''       )

								# get prompt text embeddings
								lowerCamelCase__         =  self.tokenizer(
								    _lowerCAmelCase   ,padding="""max_length"""   ,max_length=self.tokenizer.model_max_length   ,return_tensors="""pt"""   ,)
								lowerCamelCase__         =  text_inputs.input_ids

								if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
														lowerCamelCase__         =  self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]       )
														logger.warning(
														    """The following part of your input was truncated because CLIP can only handle sequences up to"""
														    F''' {self.tokenizer.model_max_length} tokens: {removed_text}'''       )
														lowerCamelCase__         =  text_input_ids[:, : self.tokenizer.model_max_length]

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

								# duplicate text embeddings for each generation per prompt, using mps friendly method
								lowerCamelCase__       ,			lowerCamelCase__       ,			lowerCamelCase__         =  text_embeddings.shape
								lowerCamelCase__         =  text_embeddings.repeat(1   ,_lowerCAmelCase   ,1       )
								lowerCamelCase__         =  text_embeddings.view(bs_embed * num_images_per_prompt   ,_lowerCAmelCase   ,-1       )

								# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
								# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
								# corresponds to doing no classifier free guidance.
								lowerCamelCase__         =  guidance_scale > 1.0
								# get unconditional embeddings for classifier free guidance
								if do_classifier_free_guidance:
														lowerCamelCase__         =  42
														if negative_prompt is None:
																				lowerCamelCase__         =  [""""""]
														elif type(_lowerCAmelCase       ) is not type(_lowerCAmelCase       ):
																				raise TypeError(
																				    F'''`negative_prompt` should be the same type to `prompt`, but got {type(_lowerCAmelCase       )} !='''
																				    F''' {type(_lowerCAmelCase       )}.'''       )
														elif isinstance(_lowerCAmelCase   ,_lowerCAmelCase       ):
																				lowerCamelCase__         =  [negative_prompt]
														elif batch_size != len(_lowerCAmelCase       ):
																				raise ValueError(
																				    F'''`negative_prompt`: {negative_prompt} has batch size {len(_lowerCAmelCase       )}, but `prompt`:'''
																				    F''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
																				    """ the batch size of `prompt`."""       )
														else:
																				lowerCamelCase__         =  negative_prompt

														lowerCamelCase__         =  text_input_ids.shape[-1]
														lowerCamelCase__         =  self.tokenizer(
														    _lowerCAmelCase   ,padding="""max_length"""   ,max_length=_lowerCAmelCase   ,truncation=_lowerCAmelCase   ,return_tensors="""pt"""   ,)
														lowerCamelCase__         =  self.text_encoder(uncond_input.input_ids.to(self.device       )       )[0]

														# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
														lowerCamelCase__         =  uncond_embeddings.shape[1]
														lowerCamelCase__         =  uncond_embeddings.repeat(_lowerCAmelCase   ,_lowerCAmelCase   ,1       )
														lowerCamelCase__         =  uncond_embeddings.view(batch_size * num_images_per_prompt   ,_lowerCAmelCase   ,-1       )

														# For classifier free guidance, we need to do two forward passes.
														# Here we concatenate the unconditional and text embeddings into a single batch
														# to avoid doing two forward passes
														lowerCamelCase__         =  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__         =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
								lowerCamelCase__         =  (batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
								lowerCamelCase__         =  text_embeddings.dtype
								if latents is None:
														if self.device.type == "mps":
																				# randn does not exist on mps
																				lowerCamelCase__         =  torch.randn(
																				    _lowerCAmelCase   ,generator=_lowerCAmelCase   ,device="""cpu"""   ,dtype=_lowerCAmelCase       ).to(self.device       )
																				lowerCamelCase__         =  torch.randn(_lowerCAmelCase   ,generator=_lowerCAmelCase   ,device="""cpu"""   ,dtype=_lowerCAmelCase       ).to(
																				    self.device       )
														else:
																				lowerCamelCase__         =  torch.randn(
																				    _lowerCAmelCase   ,generator=_lowerCAmelCase   ,device=self.device   ,dtype=_lowerCAmelCase       )
																				lowerCamelCase__         =  torch.randn(_lowerCAmelCase   ,generator=_lowerCAmelCase   ,device=self.device   ,dtype=_lowerCAmelCase       )
								else:
														if latents_reference.shape != latents_shape:
																				raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''       )
														lowerCamelCase__         =  latents_reference.to(self.device       )
														lowerCamelCase__         =  latents.to(self.device       )

								# This is the key part of the pipeline where we
								# try to ensure that the generated images w/ the same seed
								# but different sizes actually result in similar images
								lowerCamelCase__         =  (latents_shape[3] - latents_shape_reference[3]) // 2
								lowerCamelCase__         =  (latents_shape[2] - latents_shape_reference[2]) // 2
								lowerCamelCase__         =  latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
								lowerCamelCase__         =  latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
								lowerCamelCase__         =  0 if dx < 0 else dx
								lowerCamelCase__         =  0 if dy < 0 else dy
								lowerCamelCase__         =  max(-dx   ,0       )
								lowerCamelCase__         =  max(-dy   ,0       )
								# import pdb
								# pdb.set_trace()
								lowerCamelCase__         =  latents_reference[:, :, dy : dy + h, dx : dx + w]

								# set timesteps
								self.scheduler.set_timesteps(_lowerCAmelCase       )

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

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

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

														# predict the noise residual
														lowerCamelCase__         =  self.unet(_lowerCAmelCase   ,_lowerCAmelCase   ,encoder_hidden_states=_lowerCAmelCase       ).sample

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

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

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

								lowerCamelCase__         =  1 / 0.1_8215 * latents
								lowerCamelCase__         =  self.vae.decode(_lowerCAmelCase       ).sample

								lowerCamelCase__         =  (image / 2 + 0.5).clamp(0   ,1       )

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

								if self.safety_checker is not None:
														lowerCamelCase__         =  self.feature_extractor(self.numpy_to_pil(_lowerCAmelCase       )   ,return_tensors="""pt"""       ).to(
														    self.device       )
														lowerCamelCase__       ,			lowerCamelCase__         =  self.safety_checker(
														    images=_lowerCAmelCase   ,clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype       )       )
								else:
														lowerCamelCase__         =  None

								if output_type == "pil":
														lowerCamelCase__         =  self.numpy_to_pil(_lowerCAmelCase       )

								if not return_dict:
														return (image, has_nsfw_concept)

								return StableDiffusionPipelineOutput(images=_lowerCAmelCase   ,nsfw_content_detected=_lowerCAmelCase       )









'''simple docstring'''







import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


UpperCamelCase    :      Union[str, Any]  			=		logging.get_logger(__name__)
def 						A__			(				__lowerCAmelCase		:			int						):
						lowerCamelCase__         =  R"""\w+[.]\d+"""
						lowerCamelCase__         =  re.findall(__lowerCAmelCase   ,    __lowerCAmelCase						)
						for pat in pats:
												lowerCamelCase__         =  key.replace(__lowerCAmelCase   ,    """_""".join(pat.split("""."""						)						)						)
						return key
def 						A__			(				__lowerCAmelCase		:			Union[str, Any]   ,    __lowerCAmelCase		:			Union[str, Any]   ,    __lowerCAmelCase		:			Tuple						):
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""scale""",)
						if (
						    any("""norm""" in str_ for str_ in pt_tuple_key						)
						    and (pt_tuple_key[-1] == "bias")
						    and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
						    and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
						):
												lowerCamelCase__         =  pt_tuple_key[:-1] + ("""scale""",)
												return renamed_pt_tuple_key, pt_tensor
						elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
												lowerCamelCase__         =  pt_tuple_key[:-1] + ("""scale""",)
												return renamed_pt_tuple_key, pt_tensor

						# embedding
						if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
												lowerCamelCase__         =  pt_tuple_key[:-1] + ("""embedding""",)
												return renamed_pt_tuple_key, pt_tensor

						# conv layer
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""kernel""",)
						if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
												lowerCamelCase__         =  pt_tensor.transpose(2   ,    3   ,    1   ,    0						)
												return renamed_pt_tuple_key, pt_tensor

						# linear layer
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""kernel""",)
						if pt_tuple_key[-1] == "weight":
												lowerCamelCase__         =  pt_tensor.T
												return renamed_pt_tuple_key, pt_tensor

						# old PyTorch layer norm weight
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""weight""",)
						if pt_tuple_key[-1] == "gamma":
												return renamed_pt_tuple_key, pt_tensor

						# old PyTorch layer norm bias
						lowerCamelCase__         =  pt_tuple_key[:-1] + ("""bias""",)
						if pt_tuple_key[-1] == "beta":
												return renamed_pt_tuple_key, pt_tensor

						return pt_tuple_key, pt_tensor





def 						A__			(				__lowerCAmelCase		:			int   ,    __lowerCAmelCase		:			Union[str, Any]   ,    __lowerCAmelCase		:			Tuple=42						):
						# Step 1: Convert pytorch tensor to numpy
						lowerCamelCase__         =  {k: v.numpy() for k, v in pt_state_dict.items()}

						# Step 2: Since the model is stateless, get random Flax params
						lowerCamelCase__         =  flax_model.init_weights(PRNGKey(__lowerCAmelCase						)						)

						lowerCamelCase__         =  flatten_dict(__lowerCAmelCase						)
						lowerCamelCase__         =  {}

						# Need to change some parameters name to match Flax names
						for pt_key, pt_tensor in pt_state_dict.items():
												lowerCamelCase__         =  rename_key(__lowerCAmelCase						)
												lowerCamelCase__         =  tuple(renamed_pt_key.split("""."""						)						)

												# Correctly rename weight parameters
												lowerCamelCase__       ,			lowerCamelCase__         =  rename_key_and_reshape_tensor(__lowerCAmelCase   ,    __lowerCAmelCase   ,    __lowerCAmelCase						)

												if flax_key in random_flax_state_dict:
																		if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																								raise ValueError(
																								    F'''PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '''
																								    F'''{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'''						)

        # also add unexpected weight so that warning is thrown
												lowerCamelCase__         =  jnp.asarray(__lowerCAmelCase						)

						return unflatten_dict(__lowerCAmelCase						)

