






"""simple docstring"""





from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .text_encoder import MultilingualCLIP


__lowerCamelCase							=  logging.get_logger(__name__)  # pylint: disable=invalid-name

__lowerCamelCase							=  '\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/Kandinsky-2-1-prior")\n        >>> pipe_prior.to("cuda")\n\n        >>> prompt = "red cat, 4k photo"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> negative_image_emb = out.negative_image_embeds\n\n        >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1")\n        >>> pipe.to("cuda")\n\n        >>> image = pipe(\n        ...     prompt,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=negative_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ... ).images\n\n        >>> image[0].save("cat.png")\n        ```\n'



def   a    (     __UpperCAmelCase      :    Dict      ,   __UpperCAmelCase      :    Tuple      ,   __UpperCAmelCase      :    Any=8  )    ->							Optional[Any]:
						__magic_name__:	Optional[Any]							=       h // scale_factor**2
						if h % scale_factor**2 != 0:
												new_h += 1
						__magic_name__:	Dict							=       w // scale_factor**2
						if w % scale_factor**2 != 0:
												new_w += 1
						return new_h * scale_factor, new_w * scale_factor



class       __A		(   SCREAMING_SNAKE_CASE_			):

					def __init__(			self  :      Dict     ,		__snake_case  :      MultilingualCLIP     ,		__snake_case  :      XLMRobertaTokenizer     ,		__snake_case  :      UNetaDConditionModel     ,		__snake_case  :      Union[DDIMScheduler, DDPMScheduler]     ,		__snake_case  :      VQModel     ,		)			->							Optional[Any]:
											super().__init__()

											self.register_modules(
											    text_encoder=__snake_case     ,		tokenizer=__snake_case     ,		unet=__snake_case     ,		scheduler=__snake_case     ,		movq=__snake_case     ,		)
											__magic_name__:	Union[str, Any]							=       2 ** (len(self.movq.config.block_out_channels					) - 1)

					def   lowerCamelCase__   (			self  :      Dict     ,		__snake_case  :      Union[str, Any]     ,		__snake_case  :      List[str]     ,		__snake_case  :      List[str]     ,		__snake_case  :      str     ,		__snake_case  :      Any     ,		__snake_case  :      List[str]					)			->							Tuple:
											if latents is None:
																	__magic_name__:	int							=       randn_tensor(__snake_case     ,		generator=__snake_case     ,		device=__snake_case     ,		dtype=__snake_case					)
											else:
																	if latents.shape != shape:
																							raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {shape}'					)
																	__magic_name__:	List[Any]							=       latents.to(__snake_case					)

											__magic_name__:	Any							=       latents * scheduler.init_noise_sigma
											return latents

					def   lowerCamelCase__   (			self  :      Optional[Any]     ,		__snake_case  :      int     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      str     ,		__snake_case  :      Union[str, Any]     ,		__snake_case  :      Dict=None     ,		)			->							Optional[Any]:
											__magic_name__:	Any							=       len(__snake_case					) if isinstance(__snake_case     ,		__snake_case					) else 1
											# get prompt text embeddings
											__magic_name__:	Any							=       self.tokenizer(
											    __snake_case     ,		padding="""max_length"""     ,		truncation=__snake_case     ,		max_length=7_7     ,		return_attention_mask=__snake_case     ,		add_special_tokens=__snake_case     ,		return_tensors="""pt"""     ,		)

											__magic_name__:	Tuple							=       text_inputs.input_ids
											__magic_name__:	Optional[Any]							=       self.tokenizer(__snake_case     ,		padding="""longest"""     ,		return_tensors="""pt"""					).input_ids

											if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(__snake_case     ,		__snake_case					):
																	__magic_name__:	Dict							=       self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]					)
																	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}'					)

											__magic_name__:	Optional[int]							=       text_input_ids.to(__snake_case					)
											__magic_name__:	Optional[Any]							=       text_inputs.attention_mask.to(__snake_case					)

											__magic_name__,			__magic_name__:	List[str]							=       self.text_encoder(
											    input_ids=__snake_case     ,		attention_mask=__snake_case					)

											__magic_name__:	List[Any]							=       prompt_embeds.repeat_interleave(__snake_case     ,		dim=0					)
											__magic_name__:	Any							=       text_encoder_hidden_states.repeat_interleave(__snake_case     ,		dim=0					)
											__magic_name__:	Optional[int]							=       text_mask.repeat_interleave(__snake_case     ,		dim=0					)

											if do_classifier_free_guidance:
																	__magic_name__:	List[str]
																	if negative_prompt is None:
																							__magic_name__:	List[Any]							=       [""""""] * batch_size
																	elif type(__snake_case					) is not type(__snake_case					):
																							raise TypeError(
																							    F'`negative_prompt` should be the same type to `prompt`, but got {type(__snake_case					)} !='
																							    F' {type(__snake_case					)}.'					)
																	elif isinstance(__snake_case     ,		__snake_case					):
																							__magic_name__:	str							=       [negative_prompt]
																	elif batch_size != len(__snake_case					):
																							raise ValueError(
																							    F'`negative_prompt`: {negative_prompt} has batch size {len(__snake_case					)}, but `prompt`:'
																							    F' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'
																							    """ the batch size of `prompt`."""					)
																	else:
																							__magic_name__:	Any							=       negative_prompt

																	__magic_name__:	Union[str, Any]							=       self.tokenizer(
																	    __snake_case     ,		padding="""max_length"""     ,		max_length=7_7     ,		truncation=__snake_case     ,		return_attention_mask=__snake_case     ,		add_special_tokens=__snake_case     ,		return_tensors="""pt"""     ,		)
																	__magic_name__:	Optional[Any]							=       uncond_input.input_ids.to(__snake_case					)
																	__magic_name__:	Union[str, Any]							=       uncond_input.attention_mask.to(__snake_case					)

																	__magic_name__,			__magic_name__:	Any							=       self.text_encoder(
																	    input_ids=__snake_case     ,		attention_mask=__snake_case					)

																	# duplicate unconditional embeddings for each generation per prompt, using mps friendly method

																	__magic_name__:	List[Any]							=       negative_prompt_embeds.shape[1]
																	__magic_name__:	Union[str, Any]							=       negative_prompt_embeds.repeat(1     ,		__snake_case					)
																	__magic_name__:	int							=       negative_prompt_embeds.view(batch_size * num_images_per_prompt     ,		__snake_case					)

																	__magic_name__:	int							=       uncond_text_encoder_hidden_states.shape[1]
																	__magic_name__:	str							=       uncond_text_encoder_hidden_states.repeat(1     ,		__snake_case     ,		1					)
																	__magic_name__:	Optional[Any]							=       uncond_text_encoder_hidden_states.view(
																	    batch_size * num_images_per_prompt     ,		__snake_case     ,		-1					)
																	__magic_name__:	int							=       uncond_text_mask.repeat_interleave(__snake_case     ,		dim=0					)

																	# done duplicates

																	# 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
																	__magic_name__:	Any							=       torch.cat([negative_prompt_embeds, prompt_embeds]					)
																	__magic_name__:	List[Any]							=       torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]					)

																	__magic_name__:	Union[str, Any]							=       torch.cat([uncond_text_mask, text_mask]					)

											return prompt_embeds, text_encoder_hidden_states, text_mask

					def   lowerCamelCase__   (			self  :      Optional[int]     ,		__snake_case  :      List[Any]=0					)			->							Dict:
											if is_accelerate_available():
																	from accelerate import cpu_offload
											else:
																	raise ImportError("""Please install accelerate via `pip install accelerate`"""					)

											__magic_name__:	Any							=       torch.device(F'cuda:{gpu_id}'					)

											__magic_name__:	List[str]							=       [
											    self.unet,
											    self.text_encoder,
											    self.movq,
											]
											for cpu_offloaded_model in models:
																	if cpu_offloaded_model is not None:
																							cpu_offload(__snake_case     ,		__snake_case					)

					def   lowerCamelCase__   (			self  :      Tuple     ,		__snake_case  :      int=0					)			->							Optional[Any]:
											if is_accelerate_available() and is_accelerate_version(""">="""     ,		"""0.17.0.dev0"""					):
																	from accelerate import cpu_offload_with_hook
											else:
																	raise ImportError("""`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher."""					)

											__magic_name__:	str							=       torch.device(F'cuda:{gpu_id}'					)

											if self.device.type != "cpu":
																	self.to("""cpu"""     ,		silence_dtype_warnings=__snake_case					)
																	torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

											__magic_name__:	Dict							=       None
											for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
																	__magic_name__,			__magic_name__:	int							=       cpu_offload_with_hook(__snake_case     ,		__snake_case     ,		prev_module_hook=__snake_case					)

											if self.safety_checker is not None:
																	__magic_name__,			__magic_name__:	List[str]							=       cpu_offload_with_hook(self.safety_checker     ,		__snake_case     ,		prev_module_hook=__snake_case					)

											# We'll offload the last model manually.
											__magic_name__:	Dict							=       hook

					@property
					# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
					def   lowerCamelCase__   (			self  :      Dict					)			->							Optional[Any]:
											if not hasattr(self.unet     ,		"""_hf_hook"""					):
																	return self.device
											for module in self.unet.modules():
																	if (
																	    hasattr(__snake_case     ,		"""_hf_hook"""					)
																	    and hasattr(module._hf_hook     ,		"""execution_device"""					)
																	    and module._hf_hook.execution_device is not None
																	):
																							return torch.device(module._hf_hook.execution_device					)
											return self.device






					@torch.no_grad()
					@replace_example_docstring(__snake_case					)
					def __call__(			self  :      Optional[int]     ,		__snake_case  :      Union[str, List[str]]     ,		__snake_case  :      Union[torch.FloatTensor, List[torch.FloatTensor]]     ,		__snake_case  :      Union[torch.FloatTensor, List[torch.FloatTensor]]     ,		__snake_case  :      Optional[Union[str, List[str]]] = None     ,		__snake_case  :      int = 5_1_2     ,		__snake_case  :      int = 5_1_2     ,		__snake_case  :      int = 1_0_0     ,		__snake_case  :      float = 4.0     ,		__snake_case  :      int = 1     ,		__snake_case  :      Optional[Union[torch.Generator, List[torch.Generator]]] = None     ,		__snake_case  :      Optional[torch.FloatTensor] = None     ,		__snake_case  :      Optional[str] = "pil"     ,		__snake_case  :      bool = True     ,		)			->							Any:

											if isinstance(__snake_case     ,		__snake_case					):
																	__magic_name__:	Optional[int]							=       1
											elif isinstance(__snake_case     ,		__snake_case					):
																	__magic_name__:	List[Any]							=       len(__snake_case					)
											else:
																	raise ValueError(F'`prompt` has to be of type `str` or `list` but is {type(__snake_case					)}'					)

											__magic_name__:	List[Any]							=       self._execution_device

											__magic_name__:	Dict							=       batch_size * num_images_per_prompt
											__magic_name__:	Union[str, Any]							=       guidance_scale > 1.0

											__magic_name__,			__magic_name__,			__magic_name__:	str							=       self._encode_prompt(
											    __snake_case     ,		__snake_case     ,		__snake_case     ,		__snake_case     ,		__snake_case					)

											if isinstance(__snake_case     ,		__snake_case					):
																	__magic_name__:	List[str]							=       torch.cat(__snake_case     ,		dim=0					)
											if isinstance(__snake_case     ,		__snake_case					):
																	__magic_name__:	int							=       torch.cat(__snake_case     ,		dim=0					)

											if do_classifier_free_guidance:
																	__magic_name__:	Optional[Any]							=       image_embeds.repeat_interleave(__snake_case     ,		dim=0					)
																	__magic_name__:	Dict							=       negative_image_embeds.repeat_interleave(__snake_case     ,		dim=0					)

											__magic_name__:	List[Any]							=       torch.cat([negative_image_embeds, image_embeds]     ,		dim=0					).to(
											    dtype=prompt_embeds.dtype     ,		device=__snake_case					)

											self.scheduler.set_timesteps(__snake_case     ,		device=__snake_case					)
											__magic_name__:	Optional[Any]							=       self.scheduler.timesteps

											__magic_name__:	Dict							=       self.unet.config.in_channels

											__magic_name__,			__magic_name__:	Union[str, Any]							=       get_new_h_w(__snake_case     ,		__snake_case     ,		self.movq_scale_factor					)

											# create initial latent
											__magic_name__:	Any							=       self.prepare_latents(
											    (batch_size, num_channels_latents, height, width)     ,		text_encoder_hidden_states.dtype     ,		__snake_case     ,		__snake_case     ,		__snake_case     ,		self.scheduler     ,		)

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

																	__magic_name__:	Union[str, Any]							=       {"""text_embeds""": prompt_embeds, """image_embeds""": image_embeds}
																	__magic_name__:	Union[str, Any]							=       self.unet(
																	    sample=__snake_case     ,		timestep=__snake_case     ,		encoder_hidden_states=__snake_case     ,		added_cond_kwargs=__snake_case     ,		return_dict=__snake_case     ,		)[0]

																	if do_classifier_free_guidance:
																							__magic_name__,			__magic_name__:	Union[str, Any]							=       noise_pred.split(latents.shape[1]     ,		dim=1					)
																							__magic_name__,			__magic_name__:	Any							=       noise_pred.chunk(2					)
																							__magic_name__,			__magic_name__:	List[Any]							=       variance_pred.chunk(2					)
																							__magic_name__:	Dict							=       noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																							__magic_name__:	str							=       torch.cat([noise_pred, variance_pred_text]     ,		dim=1					)

																	if not (
																	    hasattr(self.scheduler.config     ,		"""variance_type"""					)
																	    and self.scheduler.config.variance_type in ["learned", "learned_range"]
																	):
																							__magic_name__,			__magic_name__:	Optional[Any]							=       noise_pred.split(latents.shape[1]     ,		dim=1					)

																	# compute the previous noisy sample x_t -> x_t-1
																	__magic_name__:	List[str]							=       self.scheduler.step(
																	    __snake_case     ,		__snake_case     ,		__snake_case     ,		generator=__snake_case     ,		).prev_sample
											# post-processing
											__magic_name__:	Dict							=       self.movq.decode(__snake_case     ,		force_not_quantize=__snake_case					)["""sample"""]

											if output_type not in ["pt", "np", "pil"]:
																	raise ValueError(F'Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}'					)

											if output_type in ["np", "pil"]:
																	__magic_name__:	Optional[int]							=       image * 0.5 + 0.5
																	__magic_name__:	Dict							=       image.clamp(0     ,		1					)
																	__magic_name__:	List[str]							=       image.cpu().permute(0     ,		2     ,		3     ,		1					).float().numpy()

											if output_type == "pil":
																	__magic_name__:	List[Any]							=       self.numpy_to_pil(__snake_case					)

											if not return_dict:
																	return (image,)

											return ImagePipelineOutput(images=__snake_case					)







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


snake_case							:       Any										=	{
    '''configuration_gpt_neo''': ['''GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''GPTNeoConfig''', '''GPTNeoOnnxConfig'''],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case							:       List[Any]										=	[
         '''GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''GPTNeoForCausalLM''',
         '''GPTNeoForQuestionAnswering''',
         '''GPTNeoForSequenceClassification''',
         '''GPTNeoForTokenClassification''',
         '''GPTNeoModel''',
         '''GPTNeoPreTrainedModel''',
         '''load_tf_weights_in_gpt_neo''',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case							:       Any										=	[
         '''FlaxGPTNeoForCausalLM''',
         '''FlaxGPTNeoModel''',
         '''FlaxGPTNeoPreTrainedModel''',
     ]


if TYPE_CHECKING:
     from .configuration_gpt_neo import GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoConfig, GPTNeoOnnxConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_gpt_neo import (
              GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST,
              GPTNeoForCausalLM,
              GPTNeoForQuestionAnswering,
              GPTNeoForSequenceClassification,
              GPTNeoForTokenClassification,
              GPTNeoModel,
              GPTNeoPreTrainedModel,
              load_tf_weights_in_gpt_neo,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel, FlaxGPTNeoPreTrainedModel


else:
     import sys

     snake_case							:       Optional[int]										=	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





