


import gc
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class      UpperCamelCase(	_a  ,				unittest.TestCase ):
					snake_case_		:					Dict    					=						KandinskyVaaControlnetPipeline
					snake_case_		:					str    					=						["""image_embeds""", """negative_image_embeds""", """hint"""]
					snake_case_		:					Optional[int]    					=						["""image_embeds""", """negative_image_embeds""", """hint"""]
					snake_case_		:					Optional[int]    					=						[
					    """generator""",
					    """height""",
					    """width""",
					    """latents""",
					    """guidance_scale""",
					    """num_inference_steps""",
					    """return_dict""",
					    """guidance_scale""",
					    """num_images_per_prompt""",
					    """output_type""",
					    """return_dict""",
					]
					snake_case_		:					Dict    					=						False


					@property
					def 				SCREAMING_SNAKE_CASE_	(    self						:		Dict    )      ->						Tuple:




											'''simple docstring'''

											return 3_2


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




											'''simple docstring'''

											return 3_2


					@property
					def 				SCREAMING_SNAKE_CASE_	(    self						:		str    )      ->						List[Any]:




											'''simple docstring'''

											return self.time_input_dim


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




											'''simple docstring'''

											return self.time_input_dim * 4


					@property
					def 				SCREAMING_SNAKE_CASE_	(    self						:		Dict    )      ->						Optional[int]:




											'''simple docstring'''

											return 1_0_0


					@property
					def 				SCREAMING_SNAKE_CASE_	(    self						:		Dict    )      ->						List[Any]:




											'''simple docstring'''

											torch.manual_seed(0    )

											__snake_case				=		{
											    "in_channels": 8,
											    # Out channels is double in channels because predicts mean and variance
											    "out_channels": 8,
											    "addition_embed_type": "image_hint",
											    "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
											    "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
											    "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
											    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
											    "layers_per_block": 1,
											    "encoder_hid_dim": self.text_embedder_hidden_size,
											    "encoder_hid_dim_type": "image_proj",
											    "cross_attention_dim": self.cross_attention_dim,
											    "attention_head_dim": 4,
											    "resnet_time_scale_shift": "scale_shift",
											    "class_embed_type": None,
											}

											__snake_case				=		UNetaDConditionModel(**SCREAMING_SNAKE_CASE    )
											return model


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




											'''simple docstring'''

											return {
											    "block_out_channels": [3_2, 3_2, 6_4, 6_4],
											    "down_block_types": [
											        "DownEncoderBlock2D",
											        "DownEncoderBlock2D",
											        "DownEncoderBlock2D",
											        "AttnDownEncoderBlock2D",
											    ],
											    "in_channels": 3,
											    "latent_channels": 4,
											    "layers_per_block": 1,
											    "norm_num_groups": 8,
											    "norm_type": "spatial",
											    "num_vq_embeddings": 1_2,
											    "out_channels": 3,
											    "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
											    "vq_embed_dim": 4,
											}


					@property
					def 				SCREAMING_SNAKE_CASE_	(    self						:		Any    )      ->						Dict:




											'''simple docstring'''

											torch.manual_seed(0    )
											__snake_case				=		VQModel(**self.dummy_movq_kwargs    )
											return model


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




											'''simple docstring'''

											__snake_case				=		self.dummy_unet
											__snake_case				=		self.dummy_movq

											__snake_case				=		DDIMScheduler(
											    num_train_timesteps=1_0_0_0  ,     beta_schedule="linear"  ,     beta_start=0.00085  ,     beta_end=0.012  ,     clip_sample=SCREAMING_SNAKE_CASE  ,     set_alpha_to_one=SCREAMING_SNAKE_CASE  ,     steps_offset=1  ,     prediction_type="epsilon"  ,     thresholding=SCREAMING_SNAKE_CASE  ,     )

											__snake_case				=		{
											    "unet": unet,
											    "scheduler": scheduler,
											    "movq": movq,
											}
											return components


					def 				SCREAMING_SNAKE_CASE_	(    self						:		int  ,     SCREAMING_SNAKE_CASE						:		int  ,     SCREAMING_SNAKE_CASE						:		Union[str, Any]=0    )      ->						Any:




											'''simple docstring'''

											__snake_case				=		floats_tensor((1, self.text_embedder_hidden_size)  ,     rng=random.Random(SCREAMING_SNAKE_CASE    )    ).to(SCREAMING_SNAKE_CASE    )
											__snake_case				=		floats_tensor((1, self.text_embedder_hidden_size)  ,     rng=random.Random(seed + 1    )    ).to(
											    SCREAMING_SNAKE_CASE    )

											# create hint
											__snake_case				=		floats_tensor((1, 3, 6_4, 6_4)  ,     rng=random.Random(SCREAMING_SNAKE_CASE    )    ).to(SCREAMING_SNAKE_CASE    )

											if str(SCREAMING_SNAKE_CASE    ).startswith("mps"    ):
																	__snake_case				=		torch.manual_seed(SCREAMING_SNAKE_CASE    )
											else:
																	__snake_case				=		torch.Generator(device=SCREAMING_SNAKE_CASE    ).manual_seed(SCREAMING_SNAKE_CASE    )
											__snake_case				=		{
											    "image_embeds": image_embeds,
											    "negative_image_embeds": negative_image_embeds,
											    "hint": hint,
											    "generator": generator,
											    "height": 6_4,
											    "width": 6_4,
											    "guidance_scale": 4.0,
											    "num_inference_steps": 2,
											    "output_type": "np",
											}
											return inputs






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




											'''simple docstring'''

											__snake_case				=		"cpu"

											__snake_case				=		self.get_dummy_components()

											__snake_case				=		self.pipeline_class(**SCREAMING_SNAKE_CASE    )
											__snake_case				=		pipe.to(SCREAMING_SNAKE_CASE    )

											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE    )

											__snake_case				=		pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE    )    )
											__snake_case				=		output.images

											__snake_case				=		pipe(
											    **self.get_dummy_inputs(SCREAMING_SNAKE_CASE    )  ,     return_dict=SCREAMING_SNAKE_CASE  ,     )[0]

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

											assert image.shape == (1, 6_4, 6_4, 3)

											__snake_case				=		np.array(
											    [0.6959826, 0.868279, 0.7558092, 0.68769467, 0.85805804, 0.65977496, 0.44885302, 0.5959111, 0.4251595]    )

											assert (
											    np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-2
											), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''

											assert (
											    np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1e-2
											), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''






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


					def 				SCREAMING_SNAKE_CASE_	(    self						:		str    )      ->						Optional[int]:




											'''simple docstring'''

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






					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple    )      ->						int:




											'''simple docstring'''

											__snake_case				=		load_numpy(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/kandinskyv22/kandinskyv22_controlnet_robotcat_fp16.npy"    )

											__snake_case				=		load_image(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/kandinskyv22/hint_image_cat.png"    )
											__snake_case				=		torch.from_numpy(np.array(SCREAMING_SNAKE_CASE    )    ).float() / 255.0
											__snake_case				=		hint.permute(2  ,     0  ,     1    ).unsqueeze(0    )

											__snake_case				=		KandinskyVaaPriorPipeline.from_pretrained(
											    "kandinsky-community/kandinsky-2-2-prior"  ,     torch_dtype=torch.floataa    )
											pipe_prior.to(SCREAMING_SNAKE_CASE    )

											__snake_case				=		KandinskyVaaControlnetPipeline.from_pretrained(
											    "kandinsky-community/kandinsky-2-2-controlnet-depth"  ,     torch_dtype=torch.floataa    )
											__snake_case				=		pipeline.to(SCREAMING_SNAKE_CASE    )
											pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE    )

											__snake_case				=		"A robot, 4k photo"

											__snake_case				=		torch.Generator(device="cuda"    ).manual_seed(0    )
											__snake_case		,       __snake_case				=		pipe_prior(
											    SCREAMING_SNAKE_CASE  ,     generator=SCREAMING_SNAKE_CASE  ,     num_inference_steps=5  ,     negative_prompt=""  ,     ).to_tuple()

											__snake_case				=		torch.Generator(device="cuda"    ).manual_seed(0    )
											__snake_case				=		pipeline(
											    image_embeds=SCREAMING_SNAKE_CASE  ,     negative_image_embeds=SCREAMING_SNAKE_CASE  ,     hint=SCREAMING_SNAKE_CASE  ,     generator=SCREAMING_SNAKE_CASE  ,     num_inference_steps=1_0_0  ,     output_type="np"  ,     )

											__snake_case				=		output.images[0]

											assert image.shape == (5_1_2, 5_1_2, 3)

											assert_mean_pixel_difference(SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE    )









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

import sentencepiece as spm

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


A				:       List[str]      =  logging.get_logger(__name__)

A				:       List[Any]      =  {'vocab_file': 'spiece.model'}

A				:       Tuple      =  {
    'vocab_file': {
        'xlnet-base-cased': 'https://huggingface.co/xlnet-base-cased/resolve/main/spiece.model',
        'xlnet-large-cased': 'https://huggingface.co/xlnet-large-cased/resolve/main/spiece.model',
    }
}

A				:       Any      =  {
    'xlnet-base-cased': None,
    'xlnet-large-cased': None,
}

# Segments (not really needed)
A				:       Tuple      =  0
A				:       str      =  1
A				:       str      =  2
A				:       Union[str, Any]      =  3
A				:       Optional[Any]      =  4






class      UpperCamelCase(	_a ):
					snake_case_		:					Union[str, Any]    					=						VOCAB_FILES_NAMES
					snake_case_		:					List[Any]    					=						PRETRAINED_VOCAB_FILES_MAP
					snake_case_		:					str    					=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					snake_case_		:					Tuple    					=						"""left"""


					def __init__(    self						:		Optional[Any]  ,     SCREAMING_SNAKE_CASE						:		Optional[int]  ,     SCREAMING_SNAKE_CASE						:		Any=False  ,     SCREAMING_SNAKE_CASE						:		List[Any]=True  ,     SCREAMING_SNAKE_CASE						:		Dict=False  ,     SCREAMING_SNAKE_CASE						:		Tuple="<s>"  ,     SCREAMING_SNAKE_CASE						:		Optional[int]="</s>"  ,     SCREAMING_SNAKE_CASE						:		List[str]="<unk>"  ,     SCREAMING_SNAKE_CASE						:		List[str]="<sep>"  ,     SCREAMING_SNAKE_CASE						:		Optional[int]="<pad>"  ,     SCREAMING_SNAKE_CASE						:		Union[str, Any]="<cls>"  ,     SCREAMING_SNAKE_CASE						:		str="<mask>"  ,     SCREAMING_SNAKE_CASE						:		int=["<eop>", "<eod>"]  ,     SCREAMING_SNAKE_CASE						:		Optional[Dict[str, Any]] = None  ,     **SCREAMING_SNAKE_CASE						:		Dict  ,     )      ->						None:




											'''simple docstring'''

											__snake_case				=		AddedToken(SCREAMING_SNAKE_CASE  ,     lstrip=SCREAMING_SNAKE_CASE  ,     rstrip=SCREAMING_SNAKE_CASE    ) if isinstance(SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE    ) else mask_token

											__snake_case				=		{} if sp_model_kwargs is None else sp_model_kwargs

											super().__init__(
											    do_lower_case=SCREAMING_SNAKE_CASE  ,     remove_space=SCREAMING_SNAKE_CASE  ,     keep_accents=SCREAMING_SNAKE_CASE  ,     bos_token=SCREAMING_SNAKE_CASE  ,     eos_token=SCREAMING_SNAKE_CASE  ,     unk_token=SCREAMING_SNAKE_CASE  ,     sep_token=SCREAMING_SNAKE_CASE  ,     pad_token=SCREAMING_SNAKE_CASE  ,     cls_token=SCREAMING_SNAKE_CASE  ,     mask_token=SCREAMING_SNAKE_CASE  ,     additional_special_tokens=SCREAMING_SNAKE_CASE  ,     sp_model_kwargs=self.sp_model_kwargs  ,     **SCREAMING_SNAKE_CASE  ,     )

											__snake_case				=		3

											__snake_case				=		do_lower_case
											__snake_case				=		remove_space
											__snake_case				=		keep_accents
											__snake_case				=		vocab_file

											__snake_case				=		spm.SentencePieceProcessor(**self.sp_model_kwargs    )
											self.sp_model.Load(SCREAMING_SNAKE_CASE    )


					@property
					def 				SCREAMING_SNAKE_CASE_	(    self						:		Optional[int]    )      ->						Union[str, Any]:




											'''simple docstring'''

											return len(self.sp_model    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		str    )      ->						Tuple:




											'''simple docstring'''

											__snake_case				=		{self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE    ): i for i in range(self.vocab_size    )}
											vocab.update(self.added_tokens_encoder    )
											return vocab


					def __getstate__(    self						:		Any    )      ->						str:




											'''simple docstring'''

											__snake_case				=		self.__dict__.copy()
											__snake_case				=		None
											return state


					def __setstate__(    self						:		List[Any]  ,     SCREAMING_SNAKE_CASE						:		Optional[int]    )      ->						int:




											'''simple docstring'''

											__snake_case				=		d

											# for backward compatibility
											if not hasattr(self  ,     "sp_model_kwargs"    ):
																	__snake_case				=		{}

											__snake_case				=		spm.SentencePieceProcessor(**self.sp_model_kwargs    )
											self.sp_model.Load(self.vocab_file    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Optional[int]  ,     SCREAMING_SNAKE_CASE						:		List[Any]    )      ->						int:




											'''simple docstring'''

											if self.remove_space:
																	__snake_case				=		" ".join(inputs.strip().split()    )
											else:
																	__snake_case				=		inputs
											__snake_case				=		outputs.replace("``"  ,     "\""    ).replace("''"  ,     "\""    )

											if not self.keep_accents:
																	__snake_case				=		unicodedata.normalize("NFKD"  ,     SCREAMING_SNAKE_CASE    )
																	__snake_case				=		"".join([c for c in outputs if not unicodedata.combining(SCREAMING_SNAKE_CASE    )]    )
											if self.do_lower_case:
																	__snake_case				=		outputs.lower()

											return outputs


					def 				SCREAMING_SNAKE_CASE_	(    self						:		int  ,     SCREAMING_SNAKE_CASE						:		str    )      ->						List[str]:




											'''simple docstring'''

											__snake_case				=		self.preprocess_text(SCREAMING_SNAKE_CASE    )
											__snake_case				=		self.sp_model.encode(SCREAMING_SNAKE_CASE  ,     out_type=SCREAMING_SNAKE_CASE    )
											__snake_case				=		[]
											for piece in pieces:
																	if len(SCREAMING_SNAKE_CASE    ) > 1 and piece[-1] == str(","    ) and piece[-2].isdigit():
																							__snake_case				=		self.sp_model.EncodeAsPieces(piece[:-1].replace(SCREAMING_SNAKE_CASE  ,     ""    )    )
																							if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
																													if len(cur_pieces[0]    ) == 1:
																																			__snake_case				=		cur_pieces[1:]
																													else:
																																			__snake_case				=		cur_pieces[0][1:]
																							cur_pieces.append(piece[-1]    )
																							new_pieces.extend(SCREAMING_SNAKE_CASE    )
																	else:
																							new_pieces.append(SCREAMING_SNAKE_CASE    )

											return new_pieces


					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		Optional[int]    )      ->						Optional[Any]:




											'''simple docstring'''

											return self.sp_model.PieceToId(SCREAMING_SNAKE_CASE    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Union[str, Any]  ,     SCREAMING_SNAKE_CASE						:		Optional[Any]    )      ->						Optional[int]:




											'''simple docstring'''

											return self.sp_model.IdToPiece(SCREAMING_SNAKE_CASE    )


					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]  ,     SCREAMING_SNAKE_CASE						:		int    )      ->						List[Any]:




											'''simple docstring'''

											__snake_case				=		"".join(SCREAMING_SNAKE_CASE    ).replace(SCREAMING_SNAKE_CASE  ,     " "    ).strip()
											return out_string


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple  ,     SCREAMING_SNAKE_CASE						:		List[int]  ,     SCREAMING_SNAKE_CASE						:		bool = False  ,     SCREAMING_SNAKE_CASE						:		bool = None  ,     SCREAMING_SNAKE_CASE						:		bool = True  ,     **SCREAMING_SNAKE_CASE						:		Optional[int]  ,     )      ->						str:




											'''simple docstring'''

											__snake_case				=		kwargs.pop("use_source_tokenizer"  ,     SCREAMING_SNAKE_CASE    )

											__snake_case				=		self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE  ,     skip_special_tokens=SCREAMING_SNAKE_CASE    )

											# To avoid mixing byte-level and unicode for byte-level BPT
											# we need to build string separately for added tokens and byte-level tokens
											# cf. https://github.com/huggingface/transformers/issues/1133
											__snake_case				=		[]
											__snake_case				=		[]
											for token in filtered_tokens:
																	if skip_special_tokens and token in self.all_special_ids:
																							continue
																	if token in self.added_tokens_encoder:
																							if current_sub_text:
																													sub_texts.append(self.convert_tokens_to_string(SCREAMING_SNAKE_CASE    )    )
																													__snake_case				=		[]
																							sub_texts.append(SCREAMING_SNAKE_CASE    )
																	else:
																							current_sub_text.append(SCREAMING_SNAKE_CASE    )
											if current_sub_text:
																	sub_texts.append(self.convert_tokens_to_string(SCREAMING_SNAKE_CASE    )    )

											# Mimic the behavior of the Rust tokenizer:
											# By default, there are no spaces between special tokens
											__snake_case				=		"".join(SCREAMING_SNAKE_CASE    )

											__snake_case				=		(
											    clean_up_tokenization_spaces
											    if clean_up_tokenization_spaces is not None
											    else self.clean_up_tokenization_spaces
											)
											if clean_up_tokenization_spaces:
																	__snake_case				=		self.clean_up_tokenization(SCREAMING_SNAKE_CASE    )
																	return clean_text
											else:
																	return text


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple  ,     SCREAMING_SNAKE_CASE						:		List[int]  ,     SCREAMING_SNAKE_CASE						:		Optional[List[int]] = None    )      ->						List[int]:




											'''simple docstring'''

											__snake_case				=		[self.sep_token_id]
											__snake_case				=		[self.cls_token_id]
											if token_ids_a is None:
																	return token_ids_a + sep + cls
											return token_ids_a + sep + token_ids_a + sep + cls


					def 				SCREAMING_SNAKE_CASE_	(    self						:		int  ,     SCREAMING_SNAKE_CASE						:		List[int]  ,     SCREAMING_SNAKE_CASE						:		Optional[List[int]] = None  ,     SCREAMING_SNAKE_CASE						:		bool = False    )      ->						List[int]:




											'''simple docstring'''

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

											if token_ids_a is not None:
																	return ([0] * len(SCREAMING_SNAKE_CASE    )) + [1] + ([0] * len(SCREAMING_SNAKE_CASE    )) + [1, 1]
											return ([0] * len(SCREAMING_SNAKE_CASE    )) + [1, 1]


					def 				SCREAMING_SNAKE_CASE_	(    self						:		Optional[int]  ,     SCREAMING_SNAKE_CASE						:		List[int]  ,     SCREAMING_SNAKE_CASE						:		Optional[List[int]] = None    )      ->						List[int]:




											'''simple docstring'''

											__snake_case				=		[self.sep_token_id]
											__snake_case				=		[2]

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






					def 				SCREAMING_SNAKE_CASE_	(    self						:		Tuple  ,     SCREAMING_SNAKE_CASE						:		str  ,     SCREAMING_SNAKE_CASE						:		Optional[str] = None    )      ->						Tuple[str]:




											'''simple docstring'''

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

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

											return (out_vocab_file,)





