





import gc
import random
import unittest

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

from diffusers import DDIMScheduler, KandinskyImgaImgPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
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       A_			(a_	,				unittest.TestCase     ):
	UpperCAmelCase__   =     KandinskyImgaImgPipeline
	UpperCAmelCase__   =     ['''prompt''', '''image_embeds''', '''negative_image_embeds''', '''image''']
	UpperCAmelCase__   =     [
	    '''prompt''',
	    '''negative_prompt''',
	    '''image_embeds''',
	    '''negative_image_embeds''',
	    '''image''',
	]
	UpperCAmelCase__   =     [
	    '''generator''',
	    '''height''',
	    '''width''',
	    '''strength''',
	    '''guidance_scale''',
	    '''negative_prompt''',
	    '''num_inference_steps''',
	    '''return_dict''',
	    '''guidance_scale''',
	    '''num_images_per_prompt''',
	    '''output_type''',
	    '''return_dict''',
	]
	UpperCAmelCase__   =     False
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				return 3_2
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				return 3_2
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				return self.time_input_dim
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				return self.time_input_dim * 4
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				return 1_0_0
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''						)
				return tokenizer
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				torch.manual_seed(0						)
				UpperCAmelCase  						=							MCLIPConfig(
				    numDims=self.cross_attention_dim				,   transformerDimensions=self.text_embedder_hidden_size				,   hidden_size=self.text_embedder_hidden_size				,   intermediate_size=3_7				,   num_attention_heads=4				,   num_hidden_layers=5				,   vocab_size=1_0_0_5				,   )

				UpperCAmelCase  						=							MultilingualCLIP(lowercase__						)
				UpperCAmelCase  						=							text_encoder.eval()

				return text_encoder
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				torch.manual_seed(0						)

				UpperCAmelCase  						=							{
				    '''in_channels''': 4,
				    # Out channels is double in channels because predicts mean and variance
				    '''out_channels''': 8,
				    '''addition_embed_type''': '''text_image''',
				    '''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''': '''text_image_proj''',
				    '''cross_attention_dim''': self.cross_attention_dim,
				    '''attention_head_dim''': 4,
				    '''resnet_time_scale_shift''': '''scale_shift''',
				    '''class_embed_type''': None,
				}

				UpperCAmelCase  						=							UNetaDConditionModel(**lowercase__						)
				return model
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				return {
				    "block_out_channels": [3_2, 6_4],
				    "down_block_types": ["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",
				    ],
				    "vq_embed_dim": 4,
				}
	@property
	def 	_lowercase							(       self						):






				'''simple docstring'''







				torch.manual_seed(0						)
				UpperCAmelCase  						=							VQModel(**self.dummy_movq_kwargs						)
				return model
	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							self.dummy_text_encoder
				UpperCAmelCase  						=							self.dummy_tokenizer
				UpperCAmelCase  						=							self.dummy_unet
				UpperCAmelCase  						=							self.dummy_movq

				UpperCAmelCase  						=							{
				    '''num_train_timesteps''': 1_0_0_0,
				    '''beta_schedule''': '''linear''',
				    '''beta_start''': 0.0_00_85,
				    '''beta_end''': 0.0_12,
				    '''clip_sample''': False,
				    '''set_alpha_to_one''': False,
				    '''steps_offset''': 0,
				    '''prediction_type''': '''epsilon''',
				    '''thresholding''': False,
				}

				UpperCAmelCase  						=							DDIMScheduler(**lowercase__						)

				UpperCAmelCase  						=							{
				    '''text_encoder''': text_encoder,
				    '''tokenizer''': tokenizer,
				    '''unet''': unet,
				    '''scheduler''': scheduler,
				    '''movq''': movq,
				}

				return components
	def 	_lowercase							(       self				,   _A				,   _A=0						):






				'''simple docstring'''







				UpperCAmelCase  						=							floats_tensor((1, self.cross_attention_dim)				,   rng=random.Random(lowercase__						)						).to(lowercase__						)
				UpperCAmelCase  						=							floats_tensor((1, self.cross_attention_dim)				,   rng=random.Random(seed + 1						)						).to(lowercase__						)
				# create init_image
				UpperCAmelCase  						=							floats_tensor((1, 3, 6_4, 6_4)				,   rng=random.Random(lowercase__						)						).to(lowercase__						)
				UpperCAmelCase  						=							image.cpu().permute(0				,   2				,   3				,   1						)[0]
				UpperCAmelCase  						=							Image.fromarray(np.uinta(lowercase__						)						).convert('''RGB'''						).resize((2_5_6, 2_5_6)						)

				if str(lowercase__						).startswith('''mps'''						):
							UpperCAmelCase  						=							torch.manual_seed(lowercase__						)
				else:
							UpperCAmelCase  						=							torch.Generator(device=lowercase__						).manual_seed(lowercase__						)
				UpperCAmelCase  						=							{
				    '''prompt''': '''horse''',
				    '''image''': init_image,
				    '''image_embeds''': image_embeds,
				    '''negative_image_embeds''': negative_image_embeds,
				    '''generator''': generator,
				    '''height''': 6_4,
				    '''width''': 6_4,
				    '''num_inference_steps''': 1_0,
				    '''guidance_scale''': 7.0,
				    '''strength''': 0.2,
				    '''output_type''': '''np''',
				}
				return inputs







	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							'''cpu'''

				UpperCAmelCase  						=							self.get_dummy_components()

				UpperCAmelCase  						=							self.pipeline_class(**lowercase__						)
				UpperCAmelCase  						=							pipe.to(lowercase__						)

				pipe.set_progress_bar_config(disable=lowercase__						)

				UpperCAmelCase  						=							pipe(**self.get_dummy_inputs(lowercase__						)						)
				UpperCAmelCase  						=							output.images

				UpperCAmelCase  						=							pipe(
				    **self.get_dummy_inputs(lowercase__						)				,   return_dict=lowercase__				,   )[0]

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

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

				UpperCAmelCase  						=							np.array(
				    [0.61_47_49_43, 0.6_07_35_39, 0.43_30_85_44, 0.5_92_82_69, 0.47_49_35_95, 0.46_75_59_73, 0.4_61_38_38, 0.45_36_87_97, 0.50_11_92_33]						)
				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       A_			(unittest.TestCase     ):
	def 	_lowercase							(       self						):






				'''simple docstring'''







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







	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							load_numpy(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
				    '''/kandinsky/kandinsky_img2img_frog.npy'''						)

				UpperCAmelCase  						=							load_image(
				    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/kandinsky/cat.png'''						)
				UpperCAmelCase  						=							'''A red cartoon frog, 4k'''

				UpperCAmelCase  						=							KandinskyPriorPipeline.from_pretrained(
				    '''kandinsky-community/kandinsky-2-1-prior'''				,   torch_dtype=torch.floataa						)
				pipe_prior.to(lowercase__						)

				UpperCAmelCase  						=							KandinskyImgaImgPipeline.from_pretrained(
				    '''kandinsky-community/kandinsky-2-1'''				,   torch_dtype=torch.floataa						)
				UpperCAmelCase  						=							pipeline.to(lowercase__						)

				pipeline.set_progress_bar_config(disable=lowercase__						)

				UpperCAmelCase  						=							torch.Generator(device='''cpu'''						).manual_seed(0						)
				UpperCAmelCase						,     UpperCAmelCase  						=							pipe_prior(
				    lowercase__				,   generator=lowercase__				,   num_inference_steps=5				,   negative_prompt=''''''				,   ).to_tuple()

				UpperCAmelCase  						=							pipeline(
				    lowercase__				,   image=lowercase__				,   image_embeds=lowercase__				,   negative_image_embeds=lowercase__				,   generator=lowercase__				,   num_inference_steps=1_0_0				,   height=7_6_8				,   width=7_6_8				,   strength=0.2				,   output_type='''np'''				,   )

				UpperCAmelCase  						=							output.images[0]

				assert image.shape == (7_6_8, 7_6_8, 3)

				assert_mean_pixel_difference(lowercase__				,   lowercase__						)




import os
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 logging


_lowercase:	int       			=		logging.get_logger(__name__)

_lowercase:	Union[str, Any]       			=		{'''vocab_file''': '''sentencepiece.bpe.model'''}

_lowercase:	Dict       			=		{
    '''vocab_file''': {
        '''camembert-base''': '''https://huggingface.co/camembert-base/resolve/main/sentencepiece.bpe.model''',
    }
}

_lowercase:	List[Any]       			=		{
    '''camembert-base''': 5_1_2,
}

_lowercase:	Dict       			=		'''▁'''



class        lowerCamelCase__				(   UpperCAmelCase				):
			UpperCamelCase__    							=VOCAB_FILES_NAMES
			UpperCamelCase__    							=PRETRAINED_VOCAB_FILES_MAP
			UpperCamelCase__    							=PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			UpperCamelCase__    							=["input_ids", "attention_mask"]





			def __init__(       self	:	str  ,  lowercase__	:	int  ,  lowercase__	:	Tuple="<s>"  ,  lowercase__	:	Optional[int]="</s>"  ,  lowercase__	:	Optional[Any]="</s>"  ,  lowercase__	:	Any="<s>"  ,  lowercase__	:	Union[str, Any]="<unk>"  ,  lowercase__	:	Union[str, Any]="<pad>"  ,  lowercase__	:	Optional[int]="<mask>"  ,  lowercase__	:	str=["<s>NOTUSED", "</s>NOTUSED"]  ,  lowercase__	:	Optional[Dict[str, Any]] = None  ,  **lowercase__	:	Union[str, Any]  ,  ):
										# Mask token behave like a normal word, i.e. include the space before it
										_lowerCAmelCase							=							AddedToken(lowercase__  ,  lstrip=lowercase__  ,  rstrip=lowercase__						) if isinstance(lowercase__  ,  lowercase__						) else mask_token

										_lowerCAmelCase							=							{} if sp_model_kwargs is None else sp_model_kwargs

										super().__init__(
										    bos_token=lowercase__  ,  eos_token=lowercase__  ,  unk_token=lowercase__  ,  sep_token=lowercase__  ,  cls_token=lowercase__  ,  pad_token=lowercase__  ,  mask_token=lowercase__  ,  additional_special_tokens=lowercase__  ,  sp_model_kwargs=self.sp_model_kwargs  ,  **lowercase__  ,  )
										_lowerCAmelCase							=							spm.SentencePieceProcessor(**self.sp_model_kwargs						)
										self.sp_model.Load(str(lowercase__						)						)
										_lowerCAmelCase							=							vocab_file
										# HACK: These tokens were added by fairseq but don't seem to be actually used when duplicated in the actual
										# sentencepiece vocabulary (this is the case for <s> and </s>
										_lowerCAmelCase							=							{'<s>NOTUSED': 0, '<pad>': 1, '</s>NOTUSED': 2, '<unk>': 3}
										_lowerCAmelCase							=							len(self.fairseq_tokens_to_ids						)
										_lowerCAmelCase							=							len(self.sp_model						) + len(self.fairseq_tokens_to_ids						)
										_lowerCAmelCase							=							{v: k for k, v in self.fairseq_tokens_to_ids.items()}





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]  ,  lowercase__	:	List[int]  ,  lowercase__	:	Optional[List[int]] = None						):

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





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any  ,  lowercase__	:	List[int]  ,  lowercase__	:	Optional[List[int]] = None  ,  lowercase__	:	bool = False						):
										if already_has_special_tokens:
																	return super().get_special_tokens_mask(
																	    token_ids_a=lowercase__  ,  token_ids_a=lowercase__  ,  already_has_special_tokens=lowercase__						)

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





			def   SCREAMING_SNAKE_CASE__     (       self	:	int  ,  lowercase__	:	List[int]  ,  lowercase__	:	Optional[List[int]] = None						):
										_lowerCAmelCase							=							[self.sep_token_id]
										_lowerCAmelCase							=							[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 + sep + token_ids_a + sep						) * [0]





			@property
			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]						):
										return len(self.fairseq_tokens_to_ids						) + len(self.sp_model						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]						):
										_lowerCAmelCase							=							{self.convert_ids_to_tokens(lowercase__						): i for i in range(self.vocab_size						)}
										vocab.update(self.added_tokens_encoder						)
										return vocab





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]  ,  lowercase__	:	str						):
										return self.sp_model.encode(lowercase__  ,  out_type=lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]  ,  lowercase__	:	Optional[Any]						):
										if token in self.fairseq_tokens_to_ids:
																	return self.fairseq_tokens_to_ids[token]
										elif self.sp_model.PieceToId(lowercase__						) == 0:
																	# Convert sentence piece unk token to fairseq unk token index
																	return self.unk_token_id
										return self.fairseq_offset + self.sp_model.PieceToId(lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]  ,  lowercase__	:	Tuple						):
										if index in self.fairseq_ids_to_tokens:
																	return self.fairseq_ids_to_tokens[index]
										return self.sp_model.IdToPiece(index - self.fairseq_offset						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	int  ,  lowercase__	:	Optional[int]						):
										_lowerCAmelCase							=							[]
										_lowerCAmelCase							=							''
										_lowerCAmelCase							=							False
										for token in tokens:
																	# make sure that special tokens are not decoded using sentencepiece model
																	if token in self.all_special_tokens:
																								if not prev_is_special:
																															out_string += " "
																								out_string += self.sp_model.decode(lowercase__						) + token
																								_lowerCAmelCase							=							True
																								_lowerCAmelCase							=							[]
																	else:
																								current_sub_tokens.append(lowercase__						)
																								_lowerCAmelCase							=							False
										out_string += self.sp_model.decode(lowercase__						)
										return out_string.strip()





			def __getstate__(       self	:	Any						):
										_lowerCAmelCase							=							self.__dict__.copy()
										_lowerCAmelCase							=							None
										return state





			def __setstate__(       self	:	Optional[Any]  ,  lowercase__	:	Union[str, Any]						):
										_lowerCAmelCase							=							d

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

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





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]  ,  lowercase__	:	str  ,  lowercase__	:	Optional[str] = None						):
										if not os.path.isdir(lowercase__						):
																	logger.error(f'Vocabulary path ({save_directory}) should be a directory'						)
																	return
										_lowerCAmelCase							=							os.path.join(
										    lowercase__  ,  (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']						)

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

										return (out_vocab_file,)






