


'''simple docstring'''





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

import torch
from transformers import CLIPTextModel, CLIPTokenizer

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin, TransformeraDModel, VQModel
from ...schedulers import VQDiffusionScheduler
from ...utils import logging
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


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







class        UpperCamelCase__   (	SCREAMING_SNAKE_CASE__							,				SCREAMING_SNAKE_CASE__   ):

					@register_to_config
					def __init__(     self       :   Union[str, Any]							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Optional[Any] = None							,					UpperCamelCase__       :   List[str] = None       ):






										'''simple docstring'''






										super().__init__()

										lowercase_  							=			learnable

										if self.learnable:
															assert hidden_size is not None, "learnable=True requires `hidden_size` to be set"
															assert length is not None, "learnable=True requires `length` to be set"

															lowercase_  							=			torch.zeros(_lowercase							,					_lowercase       )
										else:
															lowercase_  							=			None

										lowercase_  							=			torch.nn.Parameter(_lowercase       )







class        UpperCamelCase__   (	SCREAMING_SNAKE_CASE__   ):
					__SCREAMING_SNAKE_CASE       :  Union[str, Any]						=   42
					__SCREAMING_SNAKE_CASE       :  Union[str, Any]						=   42
					__SCREAMING_SNAKE_CASE       :  List[Any]						=   42
					__SCREAMING_SNAKE_CASE       :  Tuple						=   42
					__SCREAMING_SNAKE_CASE       :  Any						=   42
					__SCREAMING_SNAKE_CASE       :  List[str]						=   42

					def __init__(     self       :   int							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   List[Any]							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   List[str]							,					):






										'''simple docstring'''






										super().__init__()

										self.register_modules(
										    vqvae=_lowercase							,					transformer=_lowercase							,					text_encoder=_lowercase							,					tokenizer=_lowercase							,					scheduler=_lowercase							,					learned_classifier_free_sampling_embeddings=_lowercase							,					)

					def   UpperCAmelCase__       (     self       :   Dict							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   List[str]							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										lowercase_  							=			len(_lowercase       ) if isinstance(_lowercase							,					_lowercase       ) else 1

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

										if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
															lowercase_  							=			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}'''       )
															lowercase_  							=			text_input_ids[:, : self.tokenizer.model_max_length]
										lowercase_  							=			self.text_encoder(text_input_ids.to(self.device       )       )[0]

										# NOTE: This additional step of normalizing the text embeddings is from VQ-Diffusion.
										# While CLIP does normalize the pooled output of the text transformer when combining
										# the image and text embeddings, CLIP does not directly normalize the last hidden state.
										#
										# CLIP normalizing the pooled output.
										# https://github.com/huggingface/transformers/blob/d92e22d1f28324f513f3080e5c47c071a3916721/src/transformers/models/clip/modeling_clip.py#L1052-L1053
										lowercase_  							=			prompt_embeds / prompt_embeds.norm(dim=-1							,					keepdim=_lowercase       )

										# duplicate text embeddings for each generation per prompt
										lowercase_  							=			prompt_embeds.repeat_interleave(_lowercase							,					dim=0       )

										if do_classifier_free_guidance:
															if self.learned_classifier_free_sampling_embeddings.learnable:
																				lowercase_  							=			self.learned_classifier_free_sampling_embeddings.embeddings
																				lowercase_  							=			negative_prompt_embeds.unsqueeze(0       ).repeat(_lowercase							,					1							,					1       )
															else:
																				lowercase_  							=			[""""""] * batch_size

																				lowercase_  							=			text_input_ids.shape[-1]
																				lowercase_  							=			self.tokenizer(
																				    _lowercase							,					padding="""max_length"""							,					max_length=_lowercase							,					truncation=_lowercase							,					return_tensors="""pt"""							,					)
																				lowercase_  							=			self.text_encoder(uncond_input.input_ids.to(self.device       )       )[0]
																				# See comment for normalizing text embeddings
																				lowercase_  							=			negative_prompt_embeds / negative_prompt_embeds.norm(dim=-1							,					keepdim=_lowercase       )

															# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
															lowercase_  							=			negative_prompt_embeds.shape[1]
															lowercase_  							=			negative_prompt_embeds.repeat(1							,					_lowercase							,					1       )
															lowercase_  							=			negative_prompt_embeds.view(batch_size * num_images_per_prompt							,					_lowercase							,					-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
															lowercase_  							=			torch.cat([negative_prompt_embeds, prompt_embeds]       )

										return prompt_embeds

					@torch.no_grad()
					def __call__(     self       :   str							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   List[Any] = 100							,					UpperCamelCase__       :   Optional[int] = 5.0							,					UpperCamelCase__       :   Dict = 1.0							,					UpperCamelCase__       :   str = 1							,					UpperCamelCase__       :   Optional[int] = None							,					UpperCamelCase__       :   Optional[Any] = None							,					UpperCamelCase__       :   Any = "pil"							,					UpperCamelCase__       :   Union[str, Any] = True							,					UpperCamelCase__       :   Optional[int] = None							,					UpperCamelCase__       :   Optional[Any] = 1							,					):






										'''simple docstring'''






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

										lowercase_  							=			batch_size * num_images_per_prompt

										lowercase_  							=			guidance_scale > 1.0

										lowercase_  							=			self._encode_prompt(_lowercase							,					_lowercase							,					_lowercase       )

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

										# get the initial completely masked latents unless the user supplied it

										lowercase_  							=			(batch_size, self.transformer.num_latent_pixels)
										if latents is None:
															lowercase_  							=			self.transformer.num_vector_embeds - 1
															lowercase_  							=			torch.full(_lowercase							,					_lowercase       ).to(self.device       )
										else:
															if latents.shape != latents_shape:
																				raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}'''       )
															if (latents < 0).any() or (latents >= self.transformer.num_vector_embeds).any():
																				raise ValueError(
																				    """Unexpected latents value(s). All latents be valid embedding indices i.e. in the range 0,"""
																				    F''' {self.transformer.num_vector_embeds - 1} (inclusive).'''       )
															lowercase_  							=			latents.to(self.device       )

										# set timesteps
										self.scheduler.set_timesteps(_lowercase							,					device=self.device       )

										lowercase_  							=			self.scheduler.timesteps.to(self.device       )

										lowercase_  							=			latents

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

															# predict the un-noised image
															# model_output == `log_p_x_0`
															lowercase_  							=			self.transformer(_lowercase							,					encoder_hidden_states=_lowercase							,					timestep=_lowercase       ).sample

															if do_classifier_free_guidance:
																				lowercase_  							=			model_output.chunk(2       )
																				lowercase_  							=			model_output_uncond + guidance_scale * (model_output_text - model_output_uncond)
																				model_output -= torch.logsumexp(_lowercase							,					dim=1							,					keepdim=_lowercase       )

															lowercase_  							=			self.truncate(_lowercase							,					_lowercase       )

															# remove `log(0)`'s (`-inf`s)
															lowercase_  							=			model_output.clamp(-70       )

															# compute the previous noisy sample x_t -> x_t-1
															lowercase_  							=			self.scheduler.step(_lowercase							,					timestep=_lowercase							,					sample=_lowercase							,					generator=_lowercase       ).prev_sample

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

										lowercase_  							=			self.vqvae.config.vq_embed_dim
										lowercase_  							=			(batch_size, self.transformer.height, self.transformer.width, embedding_channels)
										lowercase_  							=			self.vqvae.quantize.get_codebook_entry(_lowercase							,					shape=_lowercase       )
										lowercase_  							=			self.vqvae.decode(_lowercase							,					force_not_quantize=_lowercase       ).sample

										lowercase_  							=			(image / 2 + 0.5).clamp(0							,					1       )
										lowercase_  							=			image.cpu().permute(0							,					2							,					3							,					1       ).numpy()

										if output_type == "pil":
															lowercase_  							=			self.numpy_to_pil(_lowercase       )

										if not return_dict:
															return (image,)

										return ImagePipelineOutput(images=_lowercase       )







					def   UpperCAmelCase__       (     self       :   Optional[int]							,					UpperCamelCase__       :   Union[str, Any]							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										lowercase_  							=			torch.sort(_lowercase							,					1							,					descending=_lowercase       )
										lowercase_  							=			torch.exp(_lowercase       )
										lowercase_  							=			sorted_p_x_0.cumsum(dim=1       ) < truncation_rate

										# Ensure that at least the largest probability is not zeroed out
										lowercase_  							=			torch.full_like(keep_mask[:, 0:1, :]							,					_lowercase       )
										lowercase_  							=			torch.cat((all_true, keep_mask)							,					dim=1       )
										lowercase_  							=			keep_mask[:, :-1, :]

										lowercase_  							=			keep_mask.gather(1							,					indices.argsort(1       )       )

										lowercase_  							=			log_p_x_0.clone()

										lowercase_  							=			-torch.inf  # -inf = log(0)

										return rv






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

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


if is_sentencepiece_available():
							from .tokenization_ta import TaTokenizer
else:
							a					  =     None


a					  =     logging.get_logger(__name__)

a					  =     {'vocab_file': 'spiece.model', 'tokenizer_file': 'tokenizer.json'}

a					  =     {
    'vocab_file': {
        't5-small': 'https://huggingface.co/t5-small/resolve/main/spiece.model',
        't5-base': 'https://huggingface.co/t5-base/resolve/main/spiece.model',
        't5-large': 'https://huggingface.co/t5-large/resolve/main/spiece.model',
        't5-3b': 'https://huggingface.co/t5-3b/resolve/main/spiece.model',
        't5-11b': 'https://huggingface.co/t5-11b/resolve/main/spiece.model',
    },
    'tokenizer_file': {
        't5-small': 'https://huggingface.co/t5-small/resolve/main/tokenizer.json',
        't5-base': 'https://huggingface.co/t5-base/resolve/main/tokenizer.json',
        't5-large': 'https://huggingface.co/t5-large/resolve/main/tokenizer.json',
        't5-3b': 'https://huggingface.co/t5-3b/resolve/main/tokenizer.json',
        't5-11b': 'https://huggingface.co/t5-11b/resolve/main/tokenizer.json',
    },
}


# TODO(PVP) - this should be removed in Transformers v5
a					  =     {
    't5-small': 5_1_2,
    't5-base': 5_1_2,
    't5-large': 5_1_2,
    't5-3b': 5_1_2,
    't5-11b': 5_1_2,
}







class        UpperCamelCase__   (	__magic_name__   ):
					__SCREAMING_SNAKE_CASE       :  Union[str, Any]						=   VOCAB_FILES_NAMES
					__SCREAMING_SNAKE_CASE       :  int						=   PRETRAINED_VOCAB_FILES_MAP
					__SCREAMING_SNAKE_CASE       :  Any						=   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					__SCREAMING_SNAKE_CASE       :  str						=   ['input_ids', 'attention_mask']
					__SCREAMING_SNAKE_CASE       :  Dict						=   TaTokenizer

					__SCREAMING_SNAKE_CASE       :  List[int]						=   []

					def __init__(     self       :   int							,					UpperCamelCase__       :   Union[str, Any]=None							,					UpperCamelCase__       :   Any=None							,					UpperCamelCase__       :   Dict="</s>"							,					UpperCamelCase__       :   Tuple="<unk>"							,					UpperCamelCase__       :   Optional[Any]="<pad>"							,					UpperCamelCase__       :   Union[str, Any]=100							,					UpperCamelCase__       :   Optional[Any]=None							,					**UpperCamelCase__       :   List[str]							,					):






										'''simple docstring'''






										if extra_ids > 0 and additional_special_tokens is None:
															lowercase_  							=			[F'''<extra_id_{i}>''' for i in range(UpperCamelCase__       )]
										elif extra_ids > 0 and additional_special_tokens is not None:
															# Check that we have the right number of extra special tokens
															lowercase_  							=			len(set(filter(lambda UpperCamelCase__       : bool("""extra_id_""" in str(UpperCamelCase__       )       )							,					UpperCamelCase__       )       )       )
															if extra_tokens != extra_ids:
																				raise ValueError(
																				    F'''Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are'''
																				    """ provided to T5Tokenizer. In this case the additional_special_tokens must include the extra_ids"""
																				    """ tokens"""       )

										super().__init__(
										    UpperCamelCase__							,					tokenizer_file=UpperCamelCase__							,					eos_token=UpperCamelCase__							,					unk_token=UpperCamelCase__							,					pad_token=UpperCamelCase__							,					extra_ids=UpperCamelCase__							,					additional_special_tokens=UpperCamelCase__							,					**UpperCamelCase__							,					)

										lowercase_  							=			vocab_file
										lowercase_  							=			False if not self.vocab_file else True
										lowercase_  							=			extra_ids

					@staticmethod
					def   UpperCAmelCase__       (     UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   int       ):






										'''simple docstring'''






										if pretrained_model_name_or_path in TaTokenizerFast.max_model_input_sizes:
															lowercase_  							=			TaTokenizerFast.max_model_input_sizes[pretrained_model_name_or_path]
															if init_max_model_length is not None and init_max_model_length != max_model_length:
																				return init_max_model_length
															elif init_max_model_length is None:
																				warnings.warn(
																				    """This tokenizer was incorrectly instantiated with a model max length of"""
																				    F''' {deprecated_max_model_length} which will be corrected in Transformers v5.\nFor now, this'''
																				    """ behavior is kept to avoid breaking backwards compatibility when padding/encoding with"""
																				    """ `truncation is True`.\n- Be aware that you SHOULD NOT rely on"""
																				    F''' {pretrained_model_name_or_path} automatically truncating your input to'''
																				    F''' {deprecated_max_model_length} when padding/encoding.\n- If you want to encode/pad to sequences'''
																				    F''' longer than {deprecated_max_model_length} you can either instantiate this tokenizer with'''
																				    """ `model_max_length` or pass `max_length` when encoding/padding.\n- To avoid this warning, please"""
																				    """ instantiate this tokenizer with `model_max_length` set to your preferred value."""							,					UpperCamelCase__							,					)

										return max_model_length

					def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   Optional[str] = None       ):






										'''simple docstring'''






										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(UpperCamelCase__       ):
															logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''       )
															return
										lowercase_  							=			os.path.join(
										    UpperCamelCase__							,					(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]       )

										if os.path.abspath(self.vocab_file       ) != os.path.abspath(UpperCamelCase__       ):
															copyfile(self.vocab_file							,					UpperCamelCase__       )
															logger.info(F'''Copy vocab file to {out_vocab_file}'''       )

										return (out_vocab_file,)

					def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   List[int]							,					UpperCamelCase__       :   Optional[List[int]] = None       ):






										'''simple docstring'''






										lowercase_  							=			token_ids_a + [self.eos_token_id]
										if token_ids_a is None:
															return self.prefix_tokens + token_ids_a
										else:
															lowercase_  							=			token_ids_a + [self.eos_token_id]
															return self.prefix_tokens + token_ids_a + token_ids_a

					def   UpperCAmelCase__       (     self       :   str							,					UpperCamelCase__       :   List[int]							,					UpperCamelCase__       :   Optional[List[int]] = None       ):






										'''simple docstring'''






										lowercase_  							=			[self.eos_token_id]

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

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										return list(
										    set(filter(lambda UpperCamelCase__       : bool(re.search(R"""<extra_id_\d+>"""							,					UpperCamelCase__       )       ) is not None							,					self.additional_special_tokens       )       )       )







					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										return [self.convert_tokens_to_ids(UpperCamelCase__       ) for token in self.get_sentinel_tokens()]


