




from typing import List, Optional, Union

import torch

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	:       List[str]		    =	logging.get_logger(__name__)  # pylint: disable=invalid-name

__lowerCAmelCase	:       List[str]		    =	'\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")\n        >>> pipe_prior.to("cuda")\n        >>> prompt = "red cat, 4k photo"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> zero_image_emb = out.negative_image_embeds\n        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")\n        >>> pipe.to("cuda")\n        >>> image = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=50,\n        ... ).images\n        >>> image[0].save("cat.png")\n        ```\n'




def     __magic_name__       (				A	:      Optional[int],						A	:      Optional[Any],						A	:      Tuple=8     ):






			'''simple docstring'''







			a					=    height // scale_factor**2
			if height % scale_factor**2 != 0:
						new_height += 1
			a					=    width // scale_factor**2
			if width % scale_factor**2 != 0:
						new_width += 1
			return new_height * scale_factor, new_width * scale_factor






class 					snake_case__							(_UpperCamelCase     ):
				"""simple docstring"""



				def __init__(   self       :  int      ,   __lowerCamelCase       :  UNetaDConditionModel      ,   __lowerCamelCase       :  DDPMScheduler      ,   __lowerCamelCase       :  VQModel      ,   )      ->  List[str]:
							super().__init__()

							self.register_modules(
							    unet=__lowerCamelCase      ,   scheduler=__lowerCamelCase      ,   movq=__lowerCamelCase      ,   )
							a					=    2 ** (len(self.movq.config.block_out_channels							) - 1)





				def 							__UpperCAmelCase							(   self       :  Any      ,   __lowerCamelCase       :  Optional[int]      ,   __lowerCamelCase       :  List[Any]      ,   __lowerCamelCase       :  Tuple      ,   __lowerCamelCase       :  Dict      ,   __lowerCamelCase       :  List[str]      ,   __lowerCamelCase       :  List[str]							)      ->  Tuple:
							if latents is None:
										a					=    randn_tensor(__lowerCamelCase      ,   generator=__lowerCamelCase      ,   device=__lowerCamelCase      ,   dtype=__lowerCamelCase							)
							else:
										if latents.shape != shape:
													raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {shape}"""							)
										a					=    latents.to(__lowerCamelCase							)

							a					=    latents * scheduler.init_noise_sigma
							return latents





				def 							__UpperCAmelCase							(   self       :  Dict      ,   __lowerCamelCase       :  Union[str, Any]=0							)      ->  List[Any]:
							if is_accelerate_available():
										from accelerate import cpu_offload
							else:
										raise ImportError("Please install accelerate via `pip install accelerate`"							)

							a					=    torch.device(f"""cuda:{gpu_id}"""							)

							a					=    [
							    self.unet,
							    self.movq,
							]
							for cpu_offloaded_model in models:
										if cpu_offloaded_model is not None:
													cpu_offload(__lowerCamelCase      ,   __lowerCamelCase							)





				def 							__UpperCAmelCase							(   self       :  Optional[int]      ,   __lowerCamelCase       :  Any=0							)      ->  str:
							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."							)

							a					=    torch.device(f"""cuda:{gpu_id}"""							)

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

							a					=    None
							for cpu_offloaded_model in [self.unet, self.movq]:
										a					,	a					=    cpu_offload_with_hook(__lowerCamelCase      ,   __lowerCamelCase      ,   prev_module_hook=__lowerCamelCase							)

							# We'll offload the last model manually.
							a					=    hook





				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def 							__UpperCAmelCase							(   self       :  List[str]							)      ->  List[str]:
							if not hasattr(self.unet      ,   "_hf_hook"							):
										return self.device
							for module in self.unet.modules():
										if (
										    hasattr(__lowerCamelCase      ,   "_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(__lowerCamelCase							)
				def __call__(   self       :  Union[str, Any]      ,   __lowerCamelCase       :  Union[torch.FloatTensor, List[torch.FloatTensor]]      ,   __lowerCamelCase       :  Union[torch.FloatTensor, List[torch.FloatTensor]]      ,   __lowerCamelCase       :  int = 5_12      ,   __lowerCamelCase       :  int = 5_12      ,   __lowerCamelCase       :  int = 1_00      ,   __lowerCamelCase       :  float = 4.0      ,   __lowerCamelCase       :  int = 1      ,   __lowerCamelCase       :  Optional[Union[torch.Generator, List[torch.Generator]]] = None      ,   __lowerCamelCase       :  Optional[torch.FloatTensor] = None      ,   __lowerCamelCase       :  Optional[str] = "pil"      ,   __lowerCamelCase       :  bool = True      ,   )      ->  Union[str, Any]:
							a					=    self._execution_device

							a					=    guidance_scale > 1.0

							if isinstance(__lowerCamelCase      ,   __lowerCamelCase							):
										a					=    torch.cat(__lowerCamelCase      ,   dim=0							)
							a					=    image_embeds.shape[0] * num_images_per_prompt
							if isinstance(__lowerCamelCase      ,   __lowerCamelCase							):
										a					=    torch.cat(__lowerCamelCase      ,   dim=0							)

							if do_classifier_free_guidance:
										a					=    image_embeds.repeat_interleave(__lowerCamelCase      ,   dim=0							)
										a					=    negative_image_embeds.repeat_interleave(__lowerCamelCase      ,   dim=0							)

							a					=    torch.cat([negative_image_embeds, image_embeds]      ,   dim=0							).to(dtype=self.unet.dtype      ,   device=__lowerCamelCase							)

							self.scheduler.set_timesteps(__lowerCamelCase      ,   device=__lowerCamelCase							)
							a					=    self.scheduler.timesteps

							a					=    self.unet.config.in_channels

							a					,	a					=    downscale_height_and_width(__lowerCamelCase      ,   __lowerCamelCase      ,   self.movq_scale_factor							)

							# create initial latent
							a					=    self.prepare_latents(
							    (batch_size, num_channels_latents, height, width)      ,   image_embeds.dtype      ,   __lowerCamelCase      ,   __lowerCamelCase      ,   __lowerCamelCase      ,   self.scheduler      ,   )

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

										a					=    {"image_embeds": image_embeds}
										a					=    self.unet(
										    sample=__lowerCamelCase      ,   timestep=__lowerCamelCase      ,   encoder_hidden_states=__lowerCamelCase      ,   added_cond_kwargs=__lowerCamelCase      ,   return_dict=__lowerCamelCase      ,   )[0]

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

										# compute the previous noisy sample x_t -> x_t-1
										a					=    self.scheduler.step(
										    __lowerCamelCase      ,   __lowerCamelCase      ,   __lowerCamelCase      ,   generator=__lowerCamelCase      ,   )[0]
							# post-processing
							a					=    self.movq.decode(__lowerCamelCase      ,   force_not_quantize=__lowerCamelCase							)["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"]:
										a					=    image * 0.5 + 0.5
										a					=    image.clamp(0      ,   1							)
										a					=    image.cpu().permute(0      ,   2      ,   3      ,   1							).float().numpy()

							if output_type == "pil":
										a					=    self.numpy_to_pil(__lowerCamelCase							)

							if not return_dict:
										return (image,)

							return ImagePipelineOutput(images=__lowerCamelCase							)












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

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
		from .tokenization_big_bird import BigBirdTokenizer
else:
		__lowerCAmelCase	:       Optional[int]     =	None

__lowerCAmelCase	:       Optional[Any]     =	logging.get_logger(__name__)
__lowerCAmelCase	:       Optional[int]     =	{'vocab_file': 'spiece.model', 'tokenizer_file': 'tokenizer.json'}

__lowerCAmelCase	:       List[Any]     =	{
    'vocab_file': {
        'google/bigbird-roberta-base': 'https://huggingface.co/google/bigbird-roberta-base/resolve/main/spiece.model',
        'google/bigbird-roberta-large': (
            'https://huggingface.co/google/bigbird-roberta-large/resolve/main/spiece.model'
        ),
        'google/bigbird-base-trivia-itc': (
            'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/spiece.model'
        ),
    },
    'tokenizer_file': {
        'google/bigbird-roberta-base': (
            'https://huggingface.co/google/bigbird-roberta-base/resolve/main/tokenizer.json'
        ),
        'google/bigbird-roberta-large': (
            'https://huggingface.co/google/bigbird-roberta-large/resolve/main/tokenizer.json'
        ),
        'google/bigbird-base-trivia-itc': (
            'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/tokenizer.json'
        ),
    },
}

__lowerCAmelCase	:       List[str]     =	{
    'google/bigbird-roberta-base': 4096,
    'google/bigbird-roberta-large': 4096,
    'google/bigbird-base-trivia-itc': 4096,
}


__lowerCAmelCase	:       Any     =	'▁'




class 					snake_case__							(_UpperCamelCase     ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_ :							Any											=      VOCAB_FILES_NAMES
				SCREAMING_SNAKE_CASE_ :							List[Any]											=      PRETRAINED_VOCAB_FILES_MAP
				SCREAMING_SNAKE_CASE_ :							Dict											=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				SCREAMING_SNAKE_CASE_ :							str											=      BigBirdTokenizer
				SCREAMING_SNAKE_CASE_ :							str											=      ["""input_ids""", """attention_mask"""]
				SCREAMING_SNAKE_CASE_ :							List[int]											=      []





				def __init__(   self       :  int      ,   __lowerCamelCase       :  Any=None      ,   __lowerCamelCase       :  List[str]=None      ,   __lowerCamelCase       :  Optional[int]="<unk>"      ,   __lowerCamelCase       :  int="<s>"      ,   __lowerCamelCase       :  Optional[Any]="</s>"      ,   __lowerCamelCase       :  Tuple="<pad>"      ,   __lowerCamelCase       :  Tuple="[SEP]"      ,   __lowerCamelCase       :  Dict="[MASK]"      ,   __lowerCamelCase       :  Tuple="[CLS]"      ,   **__lowerCamelCase       :  Optional[Any]      ,   )      ->  List[Any]:
							a					=    AddedToken(__lowerCamelCase      ,   lstrip=__lowerCamelCase      ,   rstrip=__lowerCamelCase							) if isinstance(__lowerCamelCase      ,   __lowerCamelCase							) else bos_token
							a					=    AddedToken(__lowerCamelCase      ,   lstrip=__lowerCamelCase      ,   rstrip=__lowerCamelCase							) if isinstance(__lowerCamelCase      ,   __lowerCamelCase							) else eos_token
							a					=    AddedToken(__lowerCamelCase      ,   lstrip=__lowerCamelCase      ,   rstrip=__lowerCamelCase							) if isinstance(__lowerCamelCase      ,   __lowerCamelCase							) else unk_token
							a					=    AddedToken(__lowerCamelCase      ,   lstrip=__lowerCamelCase      ,   rstrip=__lowerCamelCase							) if isinstance(__lowerCamelCase      ,   __lowerCamelCase							) else pad_token
							a					=    AddedToken(__lowerCamelCase      ,   lstrip=__lowerCamelCase      ,   rstrip=__lowerCamelCase							) if isinstance(__lowerCamelCase      ,   __lowerCamelCase							) else cls_token
							a					=    AddedToken(__lowerCamelCase      ,   lstrip=__lowerCamelCase      ,   rstrip=__lowerCamelCase							) if isinstance(__lowerCamelCase      ,   __lowerCamelCase							) else sep_token

							# Mask token behave like a normal word, i.e. include the space before it
							a					=    AddedToken(__lowerCamelCase      ,   lstrip=__lowerCamelCase      ,   rstrip=__lowerCamelCase							) if isinstance(__lowerCamelCase      ,   __lowerCamelCase							) else mask_token

							super().__init__(
							    __lowerCamelCase      ,   tokenizer_file=__lowerCamelCase      ,   bos_token=__lowerCamelCase      ,   eos_token=__lowerCamelCase      ,   unk_token=__lowerCamelCase      ,   sep_token=__lowerCamelCase      ,   pad_token=__lowerCamelCase      ,   cls_token=__lowerCamelCase      ,   mask_token=__lowerCamelCase      ,   **__lowerCamelCase      ,   )

							a					=    vocab_file
							a					=    False if not self.vocab_file else True





				def 							__UpperCAmelCase							(   self       :  str      ,   __lowerCamelCase       :  List[int]      ,   __lowerCamelCase       :  Optional[List[int]] = None							)      ->  List[int]:
							a					=    [self.sep_token_id]
							a					=    [self.cls_token_id]
							if token_ids_a is None:
										return cls + token_ids_a + sep
							return cls + token_ids_a + sep + token_ids_a + sep





				def 							__UpperCAmelCase							(   self       :  List[str]      ,   __lowerCamelCase       :  List[int]      ,   __lowerCamelCase       :  Optional[List[int]] = None      ,   __lowerCamelCase       :  bool = False							)      ->  List[int]:

							if already_has_special_tokens:
										if token_ids_a is not None:
													raise ValueError(
													    "You should not supply a second sequence if the provided sequence of "
													    "ids is already formatted with special tokens for the model."							)
										return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

							if token_ids_a is None:
										return [1] + ([0] * len(__lowerCamelCase							)) + [1]
							return [1] + ([0] * len(__lowerCamelCase							)) + [1] + ([0] * len(__lowerCamelCase							)) + [1]





				def 							__UpperCAmelCase							(   self       :  Optional[int]      ,   __lowerCamelCase       :  List[int]      ,   __lowerCamelCase       :  Optional[List[int]] = None							)      ->  List[int]:
							a					=    [self.sep_token_id]
							a					=    [self.cls_token_id]

							if token_ids_a is None:
										return len(cls + token_ids_a + sep							) * [0]
							return len(cls + token_ids_a + sep							) * [0] + len(token_ids_a + sep							) * [1]





				def 							__UpperCAmelCase							(   self       :  Tuple      ,   __lowerCamelCase       :  str      ,   __lowerCamelCase       :  Optional[str] = None							)      ->  Tuple[str]:
							if not self.can_save_slow_tokenizer:
										raise ValueError(
										    "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
										    "tokenizer."							)

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

							if os.path.abspath(self.vocab_file							) != os.path.abspath(__lowerCamelCase							):
										copyfile(self.vocab_file      ,   __lowerCamelCase							)

							return (out_vocab_file,)






