




import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()


class 							_lowerCAmelCase		(     unittest.TestCase	):




				def     __a		( self							)    ->      Dict:
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()




				@property
				def     __a		( self							)    ->      int:
											SCREAMING_SNAKE_CASE			: int	     =1
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =3
											SCREAMING_SNAKE_CASE			: List[str]	     =(32, 32)

											SCREAMING_SNAKE_CASE			: int	     =floats_tensor((batch_size, num_channels) + sizes		,		rng=random.Random(0							)							).to(snake_case_							)
											return image




				@property
				def     __a		( self							)    ->      str:
											torch.manual_seed(0							)
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =UNetaDConditionModel(
											    block_out_channels=(32, 64)		,		layers_per_block=2		,		sample_size=32		,		in_channels=4		,		out_channels=4		,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')		,		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')		,		cross_attention_dim=32		,		)
											return model




				@property
				def     __a		( self							)    ->      Optional[int]:
											torch.manual_seed(0							)
											SCREAMING_SNAKE_CASE			: List[str]	     =AutoencoderKL(
											    block_out_channels=[32, 64]		,		in_channels=3		,		out_channels=3		,		down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']		,		up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']		,		latent_channels=4		,		)
											return model




				@property
				def     __a		( self							)    ->      int:
											torch.manual_seed(0							)
											SCREAMING_SNAKE_CASE			: int	     =CLIPTextConfig(
											    bos_token_id=0		,		eos_token_id=2		,		hidden_size=32		,		intermediate_size=37		,		layer_norm_eps=1E-0_5		,		num_attention_heads=4		,		num_hidden_layers=5		,		pad_token_id=1		,		vocab_size=1_000		,		)
											return CLIPTextModel(snake_case_							)




				@property
				def     __a		( self							)    ->      int:
											def extract(*snake_case_		,		**snake_case_							):


																		class 							_lowerCAmelCase		:




																						def __init__( self							)    ->      Union[str, Any]:
																													SCREAMING_SNAKE_CASE			: Optional[int]	     =torch.ones([0]							)







																						def     __a		( self		,		snake_case_							)    ->      Optional[int]:
																													self.pixel_values.to(snake_case_							)
																													return self







																		return Out()

											return extract




				def     __a		( self							)    ->      Optional[int]:
											SCREAMING_SNAKE_CASE			: Optional[Any]	     ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
											SCREAMING_SNAKE_CASE			: Optional[int]	     =self.dummy_cond_unet
											SCREAMING_SNAKE_CASE			: int	     =DDIMScheduler(
											    beta_start=0.0_0085		,		beta_end=0.012		,		beta_schedule='''scaled_linear'''		,		clip_sample=snake_case_		,		set_alpha_to_one=snake_case_		,		)

											SCREAMING_SNAKE_CASE			: str	     =self.dummy_vae
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =self.dummy_text_encoder
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

											# make sure here that pndm scheduler skips prk
											SCREAMING_SNAKE_CASE			: Any	     =StableDiffusionPipeline(
											    unet=snake_case_		,		scheduler=snake_case_		,		vae=snake_case_		,		text_encoder=snake_case_		,		tokenizer=snake_case_		,		safety_checker=snake_case_		,		feature_extractor=self.dummy_extractor		,		)
											SCREAMING_SNAKE_CASE			: List[Any]	     =sd_pipe.to(snake_case_							)
											sd_pipe.set_progress_bar_config(disable=snake_case_							)

											SCREAMING_SNAKE_CASE			: int	     ='''A painting of a squirrel eating a burger'''

											SCREAMING_SNAKE_CASE			: List[Any]	     =torch.Generator(device=snake_case_							).manual_seed(0							)
											SCREAMING_SNAKE_CASE			: Tuple	     =sd_pipe([prompt]		,		generator=snake_case_		,		guidance_scale=6.0		,		num_inference_steps=2		,		output_type='''np'''							)
											SCREAMING_SNAKE_CASE			: Any	     =output.images

											SCREAMING_SNAKE_CASE			: Any	     =torch.Generator(device=snake_case_							).manual_seed(0							)
											SCREAMING_SNAKE_CASE			: Any	     =sd_pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=6.0		,		num_inference_steps=2		,		output_type='''np'''		,		return_dict=snake_case_		,		)[0]

											SCREAMING_SNAKE_CASE			: Tuple	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: Dict	     =image_from_tuple[0, -3:, -3:, -1]

											assert image.shape == (1, 64, 64, 3)
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =np.array([0.5753, 0.6114, 0.5001, 0.5034, 0.5470, 0.4729, 0.4971, 0.4867, 0.4867]							)

											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     __a		( self							)    ->      Optional[int]:
											SCREAMING_SNAKE_CASE			: Any	     ='''cpu'''  # ensure determinism for the device-dependent torch.Generator
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =self.dummy_cond_unet
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =PNDMScheduler(skip_prk_steps=snake_case_							)
											SCREAMING_SNAKE_CASE			: str	     =self.dummy_vae
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =self.dummy_text_encoder
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

											# make sure here that pndm scheduler skips prk
											SCREAMING_SNAKE_CASE			: Any	     =StableDiffusionPipeline(
											    unet=snake_case_		,		scheduler=snake_case_		,		vae=snake_case_		,		text_encoder=snake_case_		,		tokenizer=snake_case_		,		safety_checker=snake_case_		,		feature_extractor=self.dummy_extractor		,		)
											SCREAMING_SNAKE_CASE			: Tuple	     =sd_pipe.to(snake_case_							)
											sd_pipe.set_progress_bar_config(disable=snake_case_							)

											SCREAMING_SNAKE_CASE			: str	     ='''A painting of a squirrel eating a burger'''
											SCREAMING_SNAKE_CASE			: Dict	     =torch.Generator(device=snake_case_							).manual_seed(0							)
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =sd_pipe([prompt]		,		generator=snake_case_		,		guidance_scale=6.0		,		num_inference_steps=2		,		output_type='''np'''							)

											SCREAMING_SNAKE_CASE			: int	     =output.images

											SCREAMING_SNAKE_CASE			: int	     =torch.Generator(device=snake_case_							).manual_seed(0							)
											SCREAMING_SNAKE_CASE			: Dict	     =sd_pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=6.0		,		num_inference_steps=2		,		output_type='''np'''		,		return_dict=snake_case_		,		)[0]

											SCREAMING_SNAKE_CASE			: Optional[Any]	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: List[str]	     =image_from_tuple[0, -3:, -3:, -1]

											assert image.shape == (1, 64, 64, 3)
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =np.array([0.5122, 0.5712, 0.4825, 0.5053, 0.5646, 0.4769, 0.5179, 0.4894, 0.4994]							)

											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     __a		( self							)    ->      int:
											SCREAMING_SNAKE_CASE			: str	     =StableDiffusionPipeline.from_pretrained(
											    '''hf-internal-testing/tiny-stable-diffusion-lms-pipe'''		,		safety_checker=snake_case_							)
											assert isinstance(snake_case_		,		snake_case_							)
											assert isinstance(pipe.scheduler		,		snake_case_							)
											assert pipe.safety_checker is None

											SCREAMING_SNAKE_CASE			: Optional[int]	     =pipe('''example prompt'''		,		num_inference_steps=2							).images[0]
											assert image is not None

											# check that there's no error when saving a pipeline with one of the models being None
											with tempfile.TemporaryDirectory() as tmpdirname:
																		pipe.save_pretrained(snake_case_							)
																		SCREAMING_SNAKE_CASE			: Optional[int]	     =StableDiffusionPipeline.from_pretrained(snake_case_							)

											# sanity check that the pipeline still works
											assert pipe.safety_checker is None
											SCREAMING_SNAKE_CASE			: int	     =pipe('''example prompt'''		,		num_inference_steps=2							).images[0]
											assert image is not None







				@unittest.skipIf(torch_device != '''cuda'''		,		'''This test requires a GPU'''							)
				def     __a		( self							)    ->      str:
											SCREAMING_SNAKE_CASE			: Any	     =self.dummy_cond_unet
											SCREAMING_SNAKE_CASE			: Optional[int]	     =PNDMScheduler(skip_prk_steps=snake_case_							)
											SCREAMING_SNAKE_CASE			: Any	     =self.dummy_vae
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =self.dummy_text_encoder
											SCREAMING_SNAKE_CASE			: Tuple	     =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)

											# put models in fp16
											SCREAMING_SNAKE_CASE			: str	     =unet.half()
											SCREAMING_SNAKE_CASE			: int	     =vae.half()
											SCREAMING_SNAKE_CASE			: int	     =bert.half()

											# make sure here that pndm scheduler skips prk
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =StableDiffusionPipeline(
											    unet=snake_case_		,		scheduler=snake_case_		,		vae=snake_case_		,		text_encoder=snake_case_		,		tokenizer=snake_case_		,		safety_checker=snake_case_		,		feature_extractor=self.dummy_extractor		,		)
											SCREAMING_SNAKE_CASE			: Tuple	     =sd_pipe.to(snake_case_							)
											sd_pipe.set_progress_bar_config(disable=snake_case_							)

											SCREAMING_SNAKE_CASE			: List[str]	     ='''A painting of a squirrel eating a burger'''
											SCREAMING_SNAKE_CASE			: Tuple	     =sd_pipe([prompt]		,		num_inference_steps=2		,		output_type='''np'''							).images

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







@nightly
@require_torch_gpu
class 							_lowerCAmelCase		(     unittest.TestCase	):




				def     __a		( self							)    ->      Union[str, Any]:
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()




				def     __a		( self							)    ->      Union[str, Any]:
											SCREAMING_SNAKE_CASE			: Any	     ='''cuda'''
											SCREAMING_SNAKE_CASE			: str	     =StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''							)
											SCREAMING_SNAKE_CASE			: Dict	     =pipe.to(snake_case_							)
											pipe.set_progress_bar_config(disable=snake_case_							)

											SCREAMING_SNAKE_CASE			: Optional[Any]	     ='''a photo of a cat'''
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     ={
											    '''editing_prompt''': ['''sunglasses'''],
											    '''reverse_editing_direction''': [False],
											    '''edit_warmup_steps''': 10,
											    '''edit_guidance_scale''': 6,
											    '''edit_threshold''': 0.95,
											    '''edit_momentum_scale''': 0.5,
											    '''edit_mom_beta''': 0.6,
											}

											SCREAMING_SNAKE_CASE			: List[str]	     =3
											SCREAMING_SNAKE_CASE			: Any	     =7

											# no sega enabled
											SCREAMING_SNAKE_CASE			: List[str]	     =torch.Generator(snake_case_							)
											generator.manual_seed(snake_case_							)
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=snake_case_		,		num_inference_steps=50		,		output_type='''np'''		,		width=512		,		height=512		,		)

											SCREAMING_SNAKE_CASE			: Tuple	     =output.images
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: Any	     =[
											    0.3467_3113,
											    0.3849_2733,
											    0.3759_7352,
											    0.3408_6335,
											    0.3565_0748,
											    0.3557_9205,
											    0.338_4763,
											    0.3434_0236,
											    0.357_3271,
											]

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

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

											# with sega enabled
											# generator = torch.manual_seed(seed)
											generator.manual_seed(snake_case_							)
											SCREAMING_SNAKE_CASE			: Tuple	     =pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=snake_case_		,		num_inference_steps=50		,		output_type='''np'''		,		width=512		,		height=512		,		**snake_case_		,		)

											SCREAMING_SNAKE_CASE			: Optional[int]	     =output.images
											SCREAMING_SNAKE_CASE			: Any	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: str	     =[
											    0.4188_7826,
											    0.3772_8766,
											    0.3013_8272,
											    0.4141_6335,
											    0.4166_4985,
											    0.3628_3392,
											    0.3619_1246,
											    0.4336_4465,
											    0.4300_1732,
											]

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

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




				def     __a		( self							)    ->      Optional[Any]:
											SCREAMING_SNAKE_CASE			: Optional[Any]	     ='''cuda'''
											SCREAMING_SNAKE_CASE			: List[Any]	     =StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''							)
											SCREAMING_SNAKE_CASE			: Optional[int]	     =pipe.to(snake_case_							)
											pipe.set_progress_bar_config(disable=snake_case_							)

											SCREAMING_SNAKE_CASE			: str	     ='''an image of a crowded boulevard, realistic, 4k'''
											SCREAMING_SNAKE_CASE			: Tuple	     ={
											    '''editing_prompt''': '''crowd, crowded, people''',
											    '''reverse_editing_direction''': True,
											    '''edit_warmup_steps''': 10,
											    '''edit_guidance_scale''': 8.3,
											    '''edit_threshold''': 0.9,
											    '''edit_momentum_scale''': 0.5,
											    '''edit_mom_beta''': 0.6,
											}

											SCREAMING_SNAKE_CASE			: Tuple	     =9
											SCREAMING_SNAKE_CASE			: Optional[int]	     =7

											# no sega enabled
											SCREAMING_SNAKE_CASE			: int	     =torch.Generator(snake_case_							)
											generator.manual_seed(snake_case_							)
											SCREAMING_SNAKE_CASE			: List[Any]	     =pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=snake_case_		,		num_inference_steps=50		,		output_type='''np'''		,		width=512		,		height=512		,		)

											SCREAMING_SNAKE_CASE			: List[Any]	     =output.images
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: List[str]	     =[
											    0.4349_7998,
											    0.9181_4065,
											    0.754_0739,
											    0.5558_0205,
											    0.846_7265,
											    0.538_9691,
											    0.6257_4506,
											    0.5889_7763,
											    0.5092_6757,
											]

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

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

											# with sega enabled
											# generator = torch.manual_seed(seed)
											generator.manual_seed(snake_case_							)
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=snake_case_		,		num_inference_steps=50		,		output_type='''np'''		,		width=512		,		height=512		,		**snake_case_		,		)

											SCREAMING_SNAKE_CASE			: Tuple	     =output.images
											SCREAMING_SNAKE_CASE			: str	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: List[Any]	     =[
											    0.308_9719,
											    0.3050_0144,
											    0.2901_6042,
											    0.3063_0964,
											    0.32_5687,
											    0.2941_9225,
											    0.290_8091,
											    0.2872_3598,
											    0.2769_6294,
											]

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

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




				def     __a		( self							)    ->      int:
											SCREAMING_SNAKE_CASE			: List[str]	     ='''cuda'''
											SCREAMING_SNAKE_CASE			: Any	     =StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''							)
											SCREAMING_SNAKE_CASE			: Dict	     =pipe.to(snake_case_							)
											pipe.set_progress_bar_config(disable=snake_case_							)

											SCREAMING_SNAKE_CASE			: Any	     ='''a castle next to a river'''
											SCREAMING_SNAKE_CASE			: Optional[Any]	     ={
											    '''editing_prompt''': ['''boat on a river, boat''', '''monet, impression, sunrise'''],
											    '''reverse_editing_direction''': False,
											    '''edit_warmup_steps''': [15, 18],
											    '''edit_guidance_scale''': 6,
											    '''edit_threshold''': [0.9, 0.8],
											    '''edit_momentum_scale''': 0.5,
											    '''edit_mom_beta''': 0.6,
											}

											SCREAMING_SNAKE_CASE			: Optional[Any]	     =48
											SCREAMING_SNAKE_CASE			: List[Any]	     =7

											# no sega enabled
											SCREAMING_SNAKE_CASE			: int	     =torch.Generator(snake_case_							)
											generator.manual_seed(snake_case_							)
											SCREAMING_SNAKE_CASE			: List[Any]	     =pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=snake_case_		,		num_inference_steps=50		,		output_type='''np'''		,		width=512		,		height=512		,		)

											SCREAMING_SNAKE_CASE			: int	     =output.images
											SCREAMING_SNAKE_CASE			: str	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: int	     =[
											    0.7516_3555,
											    0.7603_7145,
											    0.6_1785,
											    0.918_9673,
											    0.862_7701,
											    0.8518_9694,
											    0.851_2813,
											    0.8701_2076,
											    0.831_2857,
											]

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

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

											# with sega enabled
											# generator = torch.manual_seed(seed)
											generator.manual_seed(snake_case_							)
											SCREAMING_SNAKE_CASE			: Optional[int]	     =pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=snake_case_		,		num_inference_steps=50		,		output_type='''np'''		,		width=512		,		height=512		,		**snake_case_		,		)

											SCREAMING_SNAKE_CASE			: Optional[int]	     =output.images
											SCREAMING_SNAKE_CASE			: Any	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: Dict	     =[
											    0.7355_3365,
											    0.753_7271,
											    0.7434_1905,
											    0.6648_0356,
											    0.647_2925,
											    0.6303_9416,
											    0.6481_2905,
											    0.674_9717,
											    0.651_7102,
											]

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

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







				def     __a		( self							)    ->      Optional[int]:
											SCREAMING_SNAKE_CASE			: Any	     ='''cuda'''
											SCREAMING_SNAKE_CASE			: List[str]	     =StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''		,		torch_dtype=torch.floataa							)
											SCREAMING_SNAKE_CASE			: List[str]	     =pipe.to(snake_case_							)
											pipe.set_progress_bar_config(disable=snake_case_							)

											SCREAMING_SNAKE_CASE			: int	     ='''a photo of a cat'''
											SCREAMING_SNAKE_CASE			: List[str]	     ={
											    '''editing_prompt''': ['''sunglasses'''],
											    '''reverse_editing_direction''': [False],
											    '''edit_warmup_steps''': 10,
											    '''edit_guidance_scale''': 6,
											    '''edit_threshold''': 0.95,
											    '''edit_momentum_scale''': 0.5,
											    '''edit_mom_beta''': 0.6,
											}

											SCREAMING_SNAKE_CASE			: str	     =3
											SCREAMING_SNAKE_CASE			: List[str]	     =7

											# no sega enabled
											SCREAMING_SNAKE_CASE			: int	     =torch.Generator(snake_case_							)
											generator.manual_seed(snake_case_							)
											SCREAMING_SNAKE_CASE			: int	     =pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=snake_case_		,		num_inference_steps=50		,		output_type='''np'''		,		width=512		,		height=512		,		)

											SCREAMING_SNAKE_CASE			: Tuple	     =output.images
											SCREAMING_SNAKE_CASE			: List[Any]	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: Optional[int]	     =[
											    0.3488_7695,
											    0.387_6953,
											    0.375,
											    0.3442_3828,
											    0.358_1543,
											    0.3571_7773,
											    0.338_3789,
											    0.3457_0312,
											    0.35_9375,
											]

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

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

											# with sega enabled
											# generator = torch.manual_seed(seed)
											generator.manual_seed(snake_case_							)
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =pipe(
											    [prompt]		,		generator=snake_case_		,		guidance_scale=snake_case_		,		num_inference_steps=50		,		output_type='''np'''		,		width=512		,		height=512		,		**snake_case_		,		)

											SCREAMING_SNAKE_CASE			: Optional[int]	     =output.images
											SCREAMING_SNAKE_CASE			: str	     =image[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =[
											    0.4228_5156,
											    0.3691_4062,
											    0.2907_7148,
											    0.4204_1016,
											    0.4191_8945,
											    0.3549_8047,
											    0.361_8164,
											    0.442_3828,
											    0.4311_5234,
											]

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

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












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

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


_A							    =	logging.get_logger(__name__)

_A							    =	{"""vocab_file""": """spiece.model""", """tokenizer_file""": """tokenizer.json"""}

_A							    =	{
    """vocab_file""": {
        """TsinghuaAI/CPM-Generate""": """https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/spiece.model""",
    },
    """tokenizer_file""": {
        """TsinghuaAI/CPM-Generate""": """https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/tokenizer.json""",
    },
}


class 							_lowerCAmelCase		(     UpperCamelCase__	):




				def __init__( self		,		snake_case_=None		,		snake_case_=None		,		snake_case_=False		,		snake_case_=True		,		snake_case_=False		,		snake_case_="<s>"		,		snake_case_="</s>"		,		snake_case_="<unk>"		,		snake_case_="<sep>"		,		snake_case_="<pad>"		,		snake_case_="<cls>"		,		snake_case_="<mask>"		,		snake_case_=["<eop>", "<eod>"]		,		**snake_case_		,		)    ->      str:
											SCREAMING_SNAKE_CASE			: Optional[int]	     =AddedToken(snake_case_		,		lstrip=snake_case_		,		rstrip=snake_case_							) if isinstance(snake_case_		,		snake_case_							) else mask_token

											super().__init__(
											    vocab_file=snake_case_		,		tokenizer_file=snake_case_		,		do_lower_case=snake_case_		,		remove_space=snake_case_		,		keep_accents=snake_case_		,		bos_token=snake_case_		,		eos_token=snake_case_		,		unk_token=snake_case_		,		sep_token=snake_case_		,		pad_token=snake_case_		,		cls_token=snake_case_		,		mask_token=snake_case_		,		additional_special_tokens=snake_case_		,		**snake_case_		,		)

											SCREAMING_SNAKE_CASE			: Dict	     =3
											SCREAMING_SNAKE_CASE			: Dict	     =do_lower_case
											SCREAMING_SNAKE_CASE			: int	     =remove_space
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =keep_accents
											SCREAMING_SNAKE_CASE			: Tuple	     =vocab_file
											SCREAMING_SNAKE_CASE			: Optional[Any]	     =False if not self.vocab_file else True

											try:
																		import jieba
											except ModuleNotFoundError as error:
																		raise error.__class__(
																		    '''You need to install jieba to use CpmTokenizer or CpmTokenizerFast. '''
																		    '''See https://pypi.org/project/jieba/ for installation.'''							)
											SCREAMING_SNAKE_CASE			: str	     =jieba
											SCREAMING_SNAKE_CASE			: Dict	     =str.maketrans(''' \n'''		,		'''\u2582\u2583'''							)




				def     __a		( self		,		snake_case_		,		snake_case_ = None							)    ->      List[int]:
											SCREAMING_SNAKE_CASE			: Dict	     =[self.sep_token_id]
											SCREAMING_SNAKE_CASE			: Any	     =[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     __a		( self		,		snake_case_		,		snake_case_ = None							)    ->      List[int]:
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =[self.sep_token_id]
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =[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     __a		( self		,		snake_case_		,		snake_case_ = 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(snake_case_							):
																		logger.error(F'Vocabulary path ({save_directory}) should be a directory'							)
																		return
											SCREAMING_SNAKE_CASE			: Union[str, Any]	     =os.path.join(
											    snake_case_		,		(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']							)

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

											return (out_vocab_file,)




				def     __a		( self		,		snake_case_		,		*snake_case_		,		**snake_case_							)    ->      Optional[Any]:
											SCREAMING_SNAKE_CASE			: int	     =[
											    ''' '''.join([x.translate(self.translator							) for x in self.jieba.cut(snake_case_		,		cut_all=snake_case_							)]							)
											    for text in batch_text_or_text_pairs
											]
											return super()._batch_encode_plus(snake_case_		,		*snake_case_		,		**snake_case_							)







				def     __a		( self		,		*snake_case_		,		**snake_case_							)    ->      str:
											SCREAMING_SNAKE_CASE			: List[Any]	     =super()._decode(*snake_case_		,		**snake_case_							)
											SCREAMING_SNAKE_CASE			: List[Any]	     =text.replace(''' '''		,		''''''							).replace('''\u2582'''		,		''' '''							).replace('''\u2583'''		,		'''\n'''							)
											return text






