






from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

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


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

lowerCamelCase       :       Union[str, Any]        =				'''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline
        >>> from diffusers.utils import load_image
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
        ... )
        >>> pipe_prior.to("cuda")

        >>> prompt = "A red cartoon frog, 4k"
        >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)

        >>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained(
        ...     "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
        ... )
        >>> pipe.to("cuda")

        >>> init_image = load_image(
        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
        ...     "/kandinsky/frog.png"
        ... )

        >>> image = pipe(
        ...     image=init_image,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ...     strength=0.2,
        ... ).images

        >>> image[0].save("red_frog.png")
        ```
'''
def 					snake_case_						(   lowerCAmelCase_				:  Tuple				,				lowerCAmelCase_				:  Dict				,				lowerCAmelCase_				:  Tuple=8					):
							__lowercase					:							List[Any]			=   height // scale_factor**2
							if height % scale_factor**2 != 0:
														new_height += 1
							__lowercase					:							List[str]			=   width // scale_factor**2
							if width % scale_factor**2 != 0:
														new_width += 1
							return new_height * scale_factor, new_width * scale_factor
def 					snake_case_						(   lowerCAmelCase_				:  Union[str, Any]				,				lowerCAmelCase_				:  Optional[int]=512				,				lowerCAmelCase_				:  Optional[Any]=512					):
							__lowercase					:							Optional[Any]			=   pil_image.resize((w, h)				,				resample=Image.BICUBIC				,				reducing_gap=1					)
							__lowercase					:							List[str]			=   np.array(pil_image.convert("""RGB"""					)					)
							__lowercase					:							Tuple			=   arr.astype(np.floataa					) / 127.5 - 1
							__lowercase					:							Any			=   np.transpose(lowerCAmelCase_				,				[2, 0, 1]					)
							__lowercase					:							Any			=   torch.from_numpy(lowerCAmelCase_					).unsqueeze(0					)
							return image



class 						lowerCAmelCase   (       __a	):

					'''simple docstring'''

					def __init__(     self :		List[Any]						,   __a :		UNetaDConditionModel						,   __a :		DDPMScheduler						,   __a :		VQModel						,   )	->     int:


												"""simple docstring"""

												super().__init__()

												self.register_modules(
												    unet=__a						,   scheduler=__a						,   movq=__a						,   )
												__lowercase					:							str			=   2 ** (len(self.movq.config.block_out_channels							) - 1)



					def 					lowerCAmelCase (     self :		Any						,   __a :		Optional[int]						,   __a :		Optional[Any]						,   __a :		Any							)	->     str:


												"""simple docstring"""

												__lowercase					:							Optional[Any]			=   min(int(num_inference_steps * strength							)						,   __a							)

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

												return timesteps, num_inference_steps - t_start



					def 					lowerCAmelCase (     self :		Tuple						,   __a :		List[Any]						,   __a :		Union[str, Any]						,   __a :		Union[str, Any]						,   __a :		Dict						,   __a :		int						,   __a :		Any						,   __a :		Optional[int]=None							)	->     List[Any]:


												"""simple docstring"""

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

												__lowercase					:							Tuple			=   image.to(device=__a						,   dtype=__a							)

												__lowercase					:							Dict			=   batch_size * num_images_per_prompt

												if image.shape[1] == 4:
																			__lowercase					:							List[str]			=   image

												else:
																			if isinstance(__a						,   __a							) and len(__a							) != batch_size:
																										raise ValueError(
																										    F"You have passed a list of generators of length {len(__a							)}, but requested an effective batch"
																										    F" size of {batch_size}. Make sure the batch size matches the length of the generators."							)

																			elif isinstance(__a						,   __a							):
																										__lowercase					:							List[Any]			=   [
																										    self.movq.encode(image[i : i + 1]							).latent_dist.sample(generator[i]							) for i in range(__a							)
																										]
																										__lowercase					:							Union[str, Any]			=   torch.cat(__a						,   dim=0							)
																			else:
																										__lowercase					:							List[str]			=   self.movq.encode(__a							).latent_dist.sample(__a							)

																			__lowercase					:							Dict			=   self.movq.config.scaling_factor * init_latents

												__lowercase					:							str			=   torch.cat([init_latents]						,   dim=0							)

												__lowercase					:							Optional[Any]			=   init_latents.shape
												__lowercase					:							Tuple			=   randn_tensor(__a						,   generator=__a						,   device=__a						,   dtype=__a							)

												# get latents
												__lowercase					:							Union[str, Any]			=   self.scheduler.add_noise(__a						,   __a						,   __a							)

												__lowercase					:							List[Any]			=   init_latents

												return latents



					def 					lowerCAmelCase (     self :		Union[str, Any]						,   __a :		int=0							)	->     List[str]:


												"""simple docstring"""

												if is_accelerate_available():
																			from accelerate import cpu_offload
												else:
																			raise ImportError("""Please install accelerate via `pip install accelerate`"""							)

												__lowercase					:							List[str]			=   torch.device(F"cuda:{gpu_id}"							)

												__lowercase					:							List[Any]			=   [
												    self.unet,
												    self.movq,
												]
												for cpu_offloaded_model in models:
																			if cpu_offloaded_model is not None:
																										cpu_offload(__a						,   __a							)



					def 					lowerCAmelCase (     self :		Optional[Any]						,   __a :		List[Any]=0							)	->     int:


												"""simple docstring"""

												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."""							)

												__lowercase					:							Any			=   torch.device(F"cuda:{gpu_id}"							)

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

												__lowercase					:							List[Any]			=   None
												for cpu_offloaded_model in [self.unet, self.movq]:
																			__lowercase	,	__lowercase					:							List[str]			=   cpu_offload_with_hook(__a						,   __a						,   prev_module_hook=__a							)

												# We'll offload the last model manually.
												__lowercase					:							Tuple			=   hook



					@property
					# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
					def 					lowerCAmelCase (     self :		Any							)	->     Optional[Any]:


												"""simple docstring"""

												if not hasattr(self.unet						,   """_hf_hook"""							):
																			return self.device
												for module in self.unet.modules():
																			if (
																			    hasattr(__a						,   """_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(__a							)
					def __call__(     self :		Optional[int]						,   __a :		Union[torch.FloatTensor, List[torch.FloatTensor]]						,   __a :		Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]]						,   __a :		Union[torch.FloatTensor, List[torch.FloatTensor]]						,   __a :		int = 512						,   __a :		int = 512						,   __a :		int = 100						,   __a :		float = 4.0						,   __a :		float = 0.3						,   __a :		int = 1						,   __a :		Optional[Union[torch.Generator, List[torch.Generator]]] = None						,   __a :		Optional[str] = "pil"						,   __a :		bool = True						,   )	->     List[str]:


												"""simple docstring"""

												__lowercase					:							List[str]			=   self._execution_device

												__lowercase					:							Any			=   guidance_scale > 1.0

												if isinstance(__a						,   __a							):
																			__lowercase					:							Union[str, Any]			=   torch.cat(__a						,   dim=0							)
												__lowercase					:							Optional[Any]			=   image_embeds.shape[0]
												if isinstance(__a						,   __a							):
																			__lowercase					:							str			=   torch.cat(__a						,   dim=0							)

												if do_classifier_free_guidance:
																			__lowercase					:							List[Any]			=   image_embeds.repeat_interleave(__a						,   dim=0							)
																			__lowercase					:							List[Any]			=   negative_image_embeds.repeat_interleave(__a						,   dim=0							)

												__lowercase					:							List[Any]			=   torch.cat([negative_image_embeds, image_embeds]						,   dim=0							).to(dtype=self.unet.dtype						,   device=__a							)

												if not isinstance(__a						,   __a							):
																			__lowercase					:							int			=   [image]
												if not all(isinstance(__a						,   (PIL.Image.Image, torch.Tensor)							) for i in image							):
																			raise ValueError(
																			    F"Input is in incorrect format: {[type(__a							) for i in image]}. Currently, we only support  PIL image and pytorch tensor"							)

												__lowercase					:							Optional[Any]			=   torch.cat([prepare_image(__a						,   __a						,   __a							) for i in image]						,   dim=0							)
												__lowercase					:							Tuple			=   image.to(dtype=image_embeds.dtype						,   device=__a							)

												__lowercase					:							Any			=   self.movq.encode(__a							)["""latents"""]
												__lowercase					:							Union[str, Any]			=   latents.repeat_interleave(__a						,   dim=0							)
												self.scheduler.set_timesteps(__a						,   device=__a							)
												__lowercase	,	__lowercase					:							Dict			=   self.get_timesteps(__a						,   __a						,   __a							)
												__lowercase					:							Union[str, Any]			=   timesteps[:1].repeat(batch_size * num_images_per_prompt							)
												__lowercase	,	__lowercase					:							List[str]			=   downscale_height_and_width(__a						,   __a						,   self.movq_scale_factor							)
												__lowercase					:							Optional[Any]			=   self.prepare_latents(
												    __a						,   __a						,   __a						,   __a						,   image_embeds.dtype						,   __a						,   __a							)
												for i, t in enumerate(self.progress_bar(__a							)							):
																			# expand the latents if we are doing classifier free guidance
																			__lowercase					:							Union[str, Any]			=   torch.cat([latents] * 2							) if do_classifier_free_guidance else latents

																			__lowercase					:							Optional[Any]			=   {"""image_embeds""": image_embeds}
																			__lowercase					:							List[Any]			=   self.unet(
																			    sample=__a						,   timestep=__a						,   encoder_hidden_states=__a						,   added_cond_kwargs=__a						,   return_dict=__a						,   )[0]

																			if do_classifier_free_guidance:
																										__lowercase	,	__lowercase					:							int			=   noise_pred.split(latents.shape[1]						,   dim=1							)
																										__lowercase	,	__lowercase					:							Dict			=   noise_pred.chunk(2							)
																										__lowercase	,	__lowercase					:							Optional[int]			=   variance_pred.chunk(2							)
																										__lowercase					:							str			=   noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																										__lowercase					:							Optional[int]			=   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"]
																			):
																										__lowercase	,	__lowercase					:							Any			=   noise_pred.split(latents.shape[1]						,   dim=1							)

																			# compute the previous noisy sample x_t -> x_t-1
																			__lowercase					:							Optional[int]			=   self.scheduler.step(
																			    __a						,   __a						,   __a						,   generator=__a						,   )[0]

												# post-processing
												__lowercase					:							List[str]			=   self.movq.decode(__a						,   force_not_quantize=__a							)["""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"]:
																			__lowercase					:							Tuple			=   image * 0.5 + 0.5
																			__lowercase					:							Optional[int]			=   image.clamp(0						,   1							)
																			__lowercase					:							str			=   image.cpu().permute(0						,   2						,   3						,   1							).float().numpy()

												if output_type == "pil":
																			__lowercase					:							Optional[Any]			=   self.numpy_to_pil(__a							)

												if not return_dict:
																			return (image,)

												return ImagePipelineOutput(images=__a							)







import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCamelCase       :       int        =				'''▁'''

lowerCamelCase       :       Optional[int]        =				{'''vocab_file''': '''spiece.model'''}

lowerCamelCase       :       str        =				{
    '''vocab_file''': {'''google/pegasus-xsum''': '''https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model'''}
}

lowerCamelCase       :       str        =				{
    '''google/pegasus-xsum''': 5_12,
}


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

class 						lowerCAmelCase   (       __a	):

					'''simple docstring'''

					_A  :	Tuple     =					VOCAB_FILES_NAMES

					_A  :	Tuple     =					VOCAB_FILES_NAMES
					_A  :	Optional[int]     =					PRETRAINED_VOCAB_FILES_MAP
					_A  :	str     =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					_A  :	Optional[int]     =					['''input_ids''', '''attention_mask''']



					def __init__(     self :		Union[str, Any]						,   __a :		int						,   __a :		Any="<pad>"						,   __a :		Optional[int]="</s>"						,   __a :		Union[str, Any]="<unk>"						,   __a :		Optional[int]="<mask_2>"						,   __a :		Optional[int]="<mask_1>"						,   __a :		Dict=None						,   __a :		List[str]=103						,   __a :		Optional[Dict[str, Any]] = None						,   **__a :		List[Any]						,   )	->     None:


												"""simple docstring"""

												__lowercase					:							Tuple			=   offset
												if additional_special_tokens is not None:
																			if not isinstance(__a						,   __a							):
																										raise TypeError(
																										    F"additional_special_tokens should be of type {type(__a							)}, but is"
																										    F" {type(__a							)}"							)

																			__lowercase					:							Dict			=   (
																			    ([mask_token_sent] + additional_special_tokens)
																			    if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
																			    else additional_special_tokens
																			)
																			# fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
																			additional_special_tokens_extended += [
																			    F"<unk_{i}>" for i in range(len(__a							)						,   self.offset - 1							)
																			]

																			if len(set(__a							)							) != len(__a							):
																										raise ValueError(
																										    """Please make sure that the provided additional_special_tokens do not contain an incorrectly"""
																										    F" shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}."							)
																			__lowercase					:							Optional[Any]			=   additional_special_tokens_extended
												else:
																			__lowercase					:							int			=   [mask_token_sent] if mask_token_sent is not None else []
																			additional_special_tokens += [F"<unk_{i}>" for i in range(2						,   self.offset							)]

												__lowercase					:							str			=   {} if sp_model_kwargs is None else sp_model_kwargs

												super().__init__(
												    eos_token=__a						,   unk_token=__a						,   mask_token=__a						,   pad_token=__a						,   mask_token_sent=__a						,   offset=__a						,   additional_special_tokens=__a						,   sp_model_kwargs=self.sp_model_kwargs						,   **__a						,   )
												__lowercase					:							Optional[Any]			=   mask_token_sent
												__lowercase					:							Dict			=   vocab_file
												__lowercase					:							Optional[Any]			=   spm.SentencePieceProcessor(**self.sp_model_kwargs							)
												self.sp_model.Load(__a							)

												# add special tokens to encoder dict
												__lowercase					:							Dict[int, str]			=   {
												    0: self.pad_token,
												    1: self.eos_token,
												}

												if self.mask_token_sent is not None:
																			self.encoder.update(
																			    {
																			        2: self.mask_token_sent,
																			        3: self.mask_token,
																			    }							)

												if self.offset > 0:
																			# entries 2-104 are only used for pretraining and called <mask_1>, <mask_2>, unk_2, ...unk_102
																			# mask_token_sent is already added to list -> so start at 1
																			self.encoder.update({i + 3: additional_special_tokens[i] for i in range(1						,   self.offset - 1							)}							)

												__lowercase					:							Dict[str, int]			=   {v: k for k, v in self.encoder.items()}



					@property
					def 					lowerCAmelCase (     self :		Dict							)	->     int:


												"""simple docstring"""

												return len(self.sp_model							) + self.offset



					def 					lowerCAmelCase (     self :		int							)	->     Dict[str, int]:


												"""simple docstring"""

												__lowercase					:							Any			=   {self.convert_ids_to_tokens(__a							): i for i in range(self.vocab_size							)}
												vocab.update(self.added_tokens_encoder							)
												return vocab



					def __getstate__(     self :		List[str]							)	->     List[str]:


												"""simple docstring"""

												__lowercase					:							Union[str, Any]			=   self.__dict__.copy()
												__lowercase					:							Optional[Any]			=   None
												return state



					def __setstate__(     self :		Tuple						,   __a :		Any							)	->     Tuple:


												"""simple docstring"""

												__lowercase					:							List[str]			=   d

												# for backward compatibility
												if not hasattr(self						,   """sp_model_kwargs"""							):
																			__lowercase					:							List[str]			=   {}

												__lowercase					:							Optional[int]			=   spm.SentencePieceProcessor(**self.sp_model_kwargs							)
												self.sp_model.Load(self.vocab_file							)



					def 					lowerCAmelCase (     self :		Dict						,   __a :		str							)	->     List[str]:


												"""simple docstring"""

												return self.sp_model.encode(__a						,   out_type=__a							)



					def 					lowerCAmelCase (     self :		List[str]						,   __a :		str							)	->     int:


												"""simple docstring"""

												if token in self.decoder:
																			return self.decoder[token]
												elif token in self.added_tokens_decoder:
																			return self.added_tokens_decoder[token]
												__lowercase					:							Union[str, Any]			=   self.sp_model.piece_to_id(__a							)
												return sp_id + self.offset



					def 					lowerCAmelCase (     self :		Dict						,   __a :		int							)	->     str:


												"""simple docstring"""

												if index in self.encoder:
																			return self.encoder[index]
												elif index in self.added_tokens_encoder:
																			return self.added_tokens_encoder[index]
												else:
																			__lowercase					:							List[Any]			=   self.sp_model.IdToPiece(index - self.offset							)
												return token



					def 					lowerCAmelCase (     self :		Union[str, Any]						,   __a :		int							)	->     Dict:


												"""simple docstring"""

												__lowercase					:							Optional[int]			=   []
												__lowercase					:							Tuple			=   """"""
												for token in tokens:
																			# make sure that special tokens are not decoded using sentencepiece model
																			if token in self.all_special_tokens:
																										out_string += self.sp_model.decode(__a							) + token
																										__lowercase					:							str			=   []
																			else:
																										current_sub_tokens.append(__a							)
												out_string += self.sp_model.decode(__a							)
												return out_string.strip()



					def 					lowerCAmelCase (     self :		int						,   __a :		Optional[Any]=False							)	->     int:


												"""simple docstring"""

												return 1



					def 					lowerCAmelCase (     self :		Optional[int]						,   __a :		List[str]							)	->     Optional[Any]:


												"""simple docstring"""

												__lowercase					:							List[Any]			=   set(self.all_special_ids							)  # call it once instead of inside list comp
												all_special_ids.remove(self.unk_token_id							)  # <unk> is only sometimes special

												return [1 if x in all_special_ids else 0 for x in seq]



					def 					lowerCAmelCase (     self :		Union[str, Any]						,   __a :		List						,   __a :		Optional[List] = None						,   __a :		bool = False							)	->     List[int]:


												"""simple docstring"""

												if already_has_special_tokens:
																			return self._special_token_mask(__a							)
												elif token_ids_a is None:
																			return self._special_token_mask(__a							) + [1]
												else:
																			return self._special_token_mask(token_ids_a + token_ids_a							) + [1]



					def 					lowerCAmelCase (     self :		Optional[int]						,   __a :		Dict						,   __a :		Tuple=None							)	->     List[int]:


												"""simple docstring"""

												if token_ids_a is None:
																			return token_ids_a + [self.eos_token_id]
												# We don't expect to process pairs, but leave the pair logic for API consistency
												return token_ids_a + token_ids_a + [self.eos_token_id]



					def 					lowerCAmelCase (     self :		Tuple						,   __a :		str						,   __a :		Optional[str] = None							)	->     Tuple[str]:


												"""simple docstring"""

												if not os.path.isdir(__a							):
																			logger.error(F"Vocabulary path ({save_directory}) should be a directory"							)
																			return
												__lowercase					:							Optional[int]			=   os.path.join(
												    __a						,   (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]							)

												if os.path.abspath(self.vocab_file							) != os.path.abspath(__a							) and os.path.isfile(self.vocab_file							):
																			copyfile(self.vocab_file						,   __a							)
												elif not os.path.isfile(self.vocab_file							):
																			with open(__a						,   """wb"""							) as fi:
																										__lowercase					:							Any			=   self.sp_model.serialized_model_proto()
																										fi.write(__a							)

												return (out_vocab_file,)