





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_lowerCamelCase 					=      False




class 	UpperCamelCase_							(  unittest.TestCase     ):


							def   _snake_case	(				self     :Dict		) -> Tuple:



									"""simple docstring"""



									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()


							@property
							def   _snake_case	(				self     :Any		) -> Optional[Any]:



									"""simple docstring"""



									return 12


							@property
							def   _snake_case	(				self     :Tuple		) -> List[Any]:



									"""simple docstring"""



									return 12


							@property
							def   _snake_case	(				self     :str		) -> Union[str, Any]:



									"""simple docstring"""



									return 32


							@property
							def   _snake_case	(				self     :List[str]		) -> Tuple:



									"""simple docstring"""



									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   VQModel(
									    block_out_channels=[32, 64]				,					in_channels=3				,					out_channels=3				,					down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]				,					up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]				,					latent_channels=3				,					num_vq_embeddings=self.num_embed				,					vq_embed_dim=3				,					)
									return model


							@property
							def   _snake_case	(				self     :Optional[int]		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)
									return tokenizer


							@property
							def   _snake_case	(				self     :Any		) -> Any:



									"""simple docstring"""



									torch.manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   CLIPTextConfig(
									    bos_token_id=0				,					eos_token_id=2				,					hidden_size=self.text_embedder_hidden_size				,					intermediate_size=37				,					layer_norm_eps=1E-05				,					num_attention_heads=4				,					num_hidden_layers=5				,					pad_token_id=1				,					vocab_size=1000				,					)
									return CLIPTextModel(__A		)


							@property
							def   _snake_case	(				self     :Optional[Any]		) -> Tuple:



									"""simple docstring"""



									torch.manual_seed(0		)

									SCREAMING_SNAKE_CASE__								=   12
									SCREAMING_SNAKE_CASE__								=   12

									SCREAMING_SNAKE_CASE__								=   {
									    """attention_bias""": True,
									    """cross_attention_dim""": 32,
									    """attention_head_dim""": height * width,
									    """num_attention_heads""": 1,
									    """num_vector_embeds""": self.num_embed,
									    """num_embeds_ada_norm""": self.num_embeds_ada_norm,
									    """norm_num_groups""": 32,
									    """sample_size""": width,
									    """activation_fn""": """geglu-approximate""",
									}

									SCREAMING_SNAKE_CASE__								=   TransformeraDModel(**__A		)
									return model


							def   _snake_case	(				self     :List[Any]		) -> int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """cpu"""

									SCREAMING_SNAKE_CASE__								=   self.dummy_vqvae
									SCREAMING_SNAKE_CASE__								=   self.dummy_text_encoder
									SCREAMING_SNAKE_CASE__								=   self.dummy_tokenizer
									SCREAMING_SNAKE_CASE__								=   self.dummy_transformer
									SCREAMING_SNAKE_CASE__								=   VQDiffusionScheduler(self.num_embed		)
									SCREAMING_SNAKE_CASE__								=   LearnedClassifierFreeSamplingEmbeddings(learnable=__A		)

									SCREAMING_SNAKE_CASE__								=   VQDiffusionPipeline(
									    vqvae=__A				,					text_encoder=__A				,					tokenizer=__A				,					transformer=__A				,					scheduler=__A				,					learned_classifier_free_sampling_embeddings=__A				,					)
									SCREAMING_SNAKE_CASE__								=   pipe.to(__A		)
									pipe.set_progress_bar_config(disable=__A		)

									SCREAMING_SNAKE_CASE__								=   """teddy bear playing in the pool"""

									SCREAMING_SNAKE_CASE__								=   torch.Generator(device=__A		).manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   pipe([prompt]				,					generator=__A				,					num_inference_steps=2				,					output_type="""np"""		)
									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   torch.Generator(device=__A		).manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   pipe(
									    [prompt]				,					generator=__A				,					output_type="""np"""				,					return_dict=__A				,					num_inference_steps=2		)[0]

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]
									SCREAMING_SNAKE_CASE__								=   image_from_tuple[0, -3:, -3:, -1]

									assert image.shape == (1, 24, 24, 3)

									SCREAMING_SNAKE_CASE__								=   np.array([0.6_5_5_1, 0.6_1_6_8, 0.5_0_0_8, 0.5_6_7_6, 0.5_6_5_9, 0.4_2_9_5, 0.6_0_7_3, 0.5_5_9_9, 0.4_9_9_2]		)

									assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
									assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2


							def   _snake_case	(				self     :List[Any]		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   """cpu"""

									SCREAMING_SNAKE_CASE__								=   self.dummy_vqvae
									SCREAMING_SNAKE_CASE__								=   self.dummy_text_encoder
									SCREAMING_SNAKE_CASE__								=   self.dummy_tokenizer
									SCREAMING_SNAKE_CASE__								=   self.dummy_transformer
									SCREAMING_SNAKE_CASE__								=   VQDiffusionScheduler(self.num_embed		)
									SCREAMING_SNAKE_CASE__								=   LearnedClassifierFreeSamplingEmbeddings(
									    learnable=__A				,					hidden_size=self.text_embedder_hidden_size				,					length=tokenizer.model_max_length		)

									SCREAMING_SNAKE_CASE__								=   VQDiffusionPipeline(
									    vqvae=__A				,					text_encoder=__A				,					tokenizer=__A				,					transformer=__A				,					scheduler=__A				,					learned_classifier_free_sampling_embeddings=__A				,					)
									SCREAMING_SNAKE_CASE__								=   pipe.to(__A		)
									pipe.set_progress_bar_config(disable=__A		)

									SCREAMING_SNAKE_CASE__								=   """teddy bear playing in the pool"""

									SCREAMING_SNAKE_CASE__								=   torch.Generator(device=__A		).manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   pipe([prompt]				,					generator=__A				,					num_inference_steps=2				,					output_type="""np"""		)
									SCREAMING_SNAKE_CASE__								=   output.images

									SCREAMING_SNAKE_CASE__								=   torch.Generator(device=__A		).manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   pipe(
									    [prompt]				,					generator=__A				,					output_type="""np"""				,					return_dict=__A				,					num_inference_steps=2		)[0]

									SCREAMING_SNAKE_CASE__								=   image[0, -3:, -3:, -1]
									SCREAMING_SNAKE_CASE__								=   image_from_tuple[0, -3:, -3:, -1]

									assert image.shape == (1, 24, 24, 3)

									SCREAMING_SNAKE_CASE__								=   np.array([0.6_6_9_3, 0.6_0_7_5, 0.4_9_5_9, 0.5_7_0_1, 0.5_5_8_3, 0.4_3_3_3, 0.6_1_7_1, 0.5_6_8_4, 0.4_9_8_8]		)

									assert np.abs(image_slice.flatten() - expected_slice		).max() < 2.0
									assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2




@slow
@require_torch_gpu
class 	UpperCamelCase_							(  unittest.TestCase     ):


							def   _snake_case	(				self     :List[Any]		) -> Optional[int]:



									"""simple docstring"""



									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()


							def   _snake_case	(				self     :List[str]		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   load_numpy(
									    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
									    """/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy"""		)

									SCREAMING_SNAKE_CASE__								=   VQDiffusionPipeline.from_pretrained("""microsoft/vq-diffusion-ithq"""		)
									SCREAMING_SNAKE_CASE__								=   pipeline.to(__A		)
									pipeline.set_progress_bar_config(disable=__A		)

									# requires GPU generator for gumbel softmax
									# don't use GPU generator in tests though
									SCREAMING_SNAKE_CASE__								=   torch.Generator(device=__A		).manual_seed(0		)
									SCREAMING_SNAKE_CASE__								=   pipeline(
									    """teddy bear playing in the pool"""				,					num_images_per_prompt=1				,					generator=__A				,					output_type="""np"""				,					)

									SCREAMING_SNAKE_CASE__								=   output.images[0]

									assert image.shape == (256, 256, 3)
									assert np.abs(expected_image - image		).max() < 2.0






from typing import Dict, List, Optional

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_lowerCamelCase 					=      logging.get_logger(__name__)


_lowerCamelCase 					=      {
    'nielsr/canine-s': 2048,
}

# Unicode defines 1,114,112 total “codepoints”
_lowerCamelCase 					=      1114112

# Below: Constants defining canonical codepoints for special, pseudo-characters.
# Copied from https://github.com/google-research/language/blob/master/language/canine/special_codepoints.py
_lowerCamelCase 					=      0

_lowerCamelCase 					=      0XE0_00
_lowerCamelCase 					=      0XE0_01
_lowerCamelCase 					=      0XE0_02
_lowerCamelCase 					=      0XE0_03
_lowerCamelCase 					=      0XE0_04

# Maps special codepoints to human-readable names.
_lowerCamelCase 					=      {
    # Special symbols are represented using codepoints values that are valid,
    # but designated as "Private Use", meaning that they will never be assigned
    # characters by the Unicode Consortium, and are thus safe for use here.
    #
    # NOTE: Do *NOT* add any sort of [UNK_CHAR] here. They are explicitly
    # excluded and should fail with a hard error.
    CLS: "[CLS]",
    SEP: "[SEP]",
    BOS: "[BOS]",
    MASK: "[MASK]",
    PAD: "[PAD]",
    RESERVED: "[RESERVED]",
}

# Maps special codepoint human-readable names to their codepoint values.
_lowerCamelCase 					=      {name: codepoint for codepoint, name in SPECIAL_CODEPOINTS.items()}




class 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES


							def __init__(				self     :str				,					__A     :str=chr(__A		)				,					__A     :str=chr(__A		)				,					__A     :Dict=chr(__A		)				,					__A     :str=chr(__A		)				,					__A     :Union[str, Any]=chr(__A		)				,					__A     :str=chr(__A		)				,					__A     :int=False				,					__A     :int=2048				,					**__A     :Dict				,					) -> Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   AddedToken(__A				,					lstrip=__A				,					rstrip=__A		) if isinstance(__A				,					__A		) else bos_token
									SCREAMING_SNAKE_CASE__								=   AddedToken(__A				,					lstrip=__A				,					rstrip=__A		) if isinstance(__A				,					__A		) else eos_token
									SCREAMING_SNAKE_CASE__								=   AddedToken(__A				,					lstrip=__A				,					rstrip=__A		) if isinstance(__A				,					__A		) else sep_token
									SCREAMING_SNAKE_CASE__								=   AddedToken(__A				,					lstrip=__A				,					rstrip=__A		) if isinstance(__A				,					__A		) else cls_token
									SCREAMING_SNAKE_CASE__								=   AddedToken(__A				,					lstrip=__A				,					rstrip=__A		) if isinstance(__A				,					__A		) else pad_token

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

									super().__init__(
									    bos_token=__A				,					eos_token=__A				,					sep_token=__A				,					cls_token=__A				,					pad_token=__A				,					mask_token=__A				,					add_prefix_space=__A				,					model_max_length=__A				,					**__A				,					)

									# Creates a mapping for looking up the IDs of special symbols.
									SCREAMING_SNAKE_CASE__								=   {}
									for codepoint, name in SPECIAL_CODEPOINTS.items():
											SCREAMING_SNAKE_CASE__								=   codepoint

									# Creates a mapping for looking up the string forms of special symbol IDs.
									SCREAMING_SNAKE_CASE__								=   {
									    codepoint: name for name, codepoint in self._special_codepoints.items()
									}

									SCREAMING_SNAKE_CASE__								=   UNICODE_VOCAB_SIZE
									SCREAMING_SNAKE_CASE__								=   len(self._special_codepoints		)


							@property
							def   _snake_case	(				self     :Optional[Any]		) -> int:



									"""simple docstring"""



									return self._unicode_vocab_size


							def   _snake_case	(				self     :Tuple				,					__A     :str		) -> List[str]:



									"""simple docstring"""



									return list(__A		)


							def   _snake_case	(				self     :Optional[Any]				,					__A     :str		) -> int:



									"""simple docstring"""



									try:
											return ord(__A		)
									except TypeError:
											raise ValueError(f'''invalid token: \'{token}\''''		)


							def   _snake_case	(				self     :str				,					__A     :int		) -> str:



									"""simple docstring"""



									try:
											if index in SPECIAL_CODEPOINTS:
													return SPECIAL_CODEPOINTS[index]
											return chr(__A		)
									except TypeError:
											raise ValueError(f'''invalid id: {index}'''		)


							def   _snake_case	(				self     :Union[str, Any]				,					__A     :Optional[int]		) -> Any:



									"""simple docstring"""



									return "".join(__A		)


							def   _snake_case	(				self     :Optional[Any]				,					__A     :List[int]				,					__A     :Optional[List[int]] = None		) -> List[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   [self.sep_token_id]
									SCREAMING_SNAKE_CASE__								=   [self.cls_token_id]

									SCREAMING_SNAKE_CASE__								=   cls + token_ids_a + sep
									if token_ids_a is not None:
											result += token_ids_a + sep
									return result


							def   _snake_case	(				self     :List[Any]				,					__A     :List[int]				,					__A     :Optional[List[int]] = None				,					__A     :bool = False		) -> List[int]:



									"""simple docstring"""



									if already_has_special_tokens:
											return super().get_special_tokens_mask(
											    token_ids_a=__A				,					token_ids_a=__A				,					already_has_special_tokens=__A		)

									SCREAMING_SNAKE_CASE__								=   [1] + ([0] * len(__A		)) + [1]
									if token_ids_a is not None:
											result += ([0] * len(__A		)) + [1]
									return result


							def   _snake_case	(				self     :List[str]				,					__A     :List[int]				,					__A     :Optional[List[int]] = None		) -> List[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   [self.sep_token_id]
									SCREAMING_SNAKE_CASE__								=   [self.cls_token_id]

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


							def   _snake_case	(				self     :int				,					__A     :str				,					__A     :Optional[str] = None		) -> Any:



									"""simple docstring"""



									return ()