



import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class  lowerCamelCase__    (							lowerCamelCase__   ,       lowerCamelCase__   ,       unittest.TestCase):




		'''simple docstring'''



		snake_case_			=StableDiffusionSAGPipeline
		snake_case_			=TEXT_TO_IMAGE_PARAMS
		snake_case_			=TEXT_TO_IMAGE_BATCH_PARAMS
		snake_case_			=TEXT_TO_IMAGE_IMAGE_PARAMS
		snake_case_			=TEXT_TO_IMAGE_IMAGE_PARAMS
		snake_case_			=False



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

						"""simple docstring"""






						torch.manual_seed(0				)
						lowerCAmelCase__  :       List[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	,)
						lowerCAmelCase__  :       Optional[Any]				=  DDIMScheduler(
						    beta_start=0.0_0085	,beta_end=0.012	,beta_schedule='''scaled_linear'''	,clip_sample=__lowerCamelCase	,set_alpha_to_one=__lowerCamelCase	,)
						torch.manual_seed(0				)
						lowerCAmelCase__  :       List[Any]				=  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	,)
						torch.manual_seed(0				)
						lowerCAmelCase__  :       Union[str, Any]				=  CLIPTextConfig(
						    bos_token_id=0	,eos_token_id=2	,hidden_size=32	,intermediate_size=37	,layer_norm_eps=1e-05	,num_attention_heads=4	,num_hidden_layers=5	,pad_token_id=1	,vocab_size=10_00	,)
						lowerCAmelCase__  :       int				=  CLIPTextModel(__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)

						lowerCAmelCase__  :       str				=  {
						    '''unet''': unet,
						    '''scheduler''': scheduler,
						    '''vae''': vae,
						    '''text_encoder''': text_encoder,
						    '''tokenizer''': tokenizer,
						    '''safety_checker''': None,
						    '''feature_extractor''': None,
						}
						return components



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase=0				)  ->    List[str]:

						"""simple docstring"""






						if str(__lowerCamelCase				).startswith('''mps'''				):
										lowerCAmelCase__  :       Union[str, Any]				=  torch.manual_seed(__lowerCamelCase				)
						else:
										lowerCAmelCase__  :       Any				=  torch.Generator(device=__lowerCamelCase				).manual_seed(__lowerCamelCase				)
						lowerCAmelCase__  :       str				=  {
						    '''prompt''': '''.''',
						    '''generator''': generator,
						    '''num_inference_steps''': 2,
						    '''guidance_scale''': 1.0,
						    '''sag_scale''': 1.0,
						    '''output_type''': '''numpy''',
						}
						return inputs



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

						"""simple docstring"""






						super().test_inference_batch_single_identical(expected_max_diff=3e-3				)



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




		'''simple docstring'''



		def 	lowerCAmelCase__   (self				)  ->    List[Any]:

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''				)
						lowerCAmelCase__  :       Optional[int]				=  sag_pipe.to(__lowerCamelCase				)
						sag_pipe.set_progress_bar_config(disable=__lowerCamelCase				)

						lowerCAmelCase__  :       Union[str, Any]				=  '''.'''
						lowerCAmelCase__  :       Dict				=  torch.manual_seed(0				)
						lowerCAmelCase__  :       Dict				=  sag_pipe(
						    [prompt]	,generator=__lowerCamelCase	,guidance_scale=7.5	,sag_scale=1.0	,num_inference_steps=20	,output_type='''np'''				)

						lowerCAmelCase__  :       List[Any]				=  output.images

						lowerCAmelCase__  :       int				=  image[0, -3:, -3:, -1]

						assert image.shape == (1, 5_12, 5_12, 3)
						lowerCAmelCase__  :       Any				=  np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949]				)

						assert np.abs(image_slice.flatten() - expected_slice				).max() < 5e-2



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
						lowerCAmelCase__  :       List[Any]				=  sag_pipe.to(__lowerCamelCase				)
						sag_pipe.set_progress_bar_config(disable=__lowerCamelCase				)

						lowerCAmelCase__  :       int				=  '''.'''
						lowerCAmelCase__  :       Tuple				=  torch.manual_seed(0				)
						lowerCAmelCase__  :       Any				=  sag_pipe(
						    [prompt]	,generator=__lowerCamelCase	,guidance_scale=7.5	,sag_scale=1.0	,num_inference_steps=20	,output_type='''np'''				)

						lowerCAmelCase__  :       Dict				=  output.images

						lowerCAmelCase__  :       Any				=  image[0, -3:, -3:, -1]

						assert image.shape == (1, 5_12, 5_12, 3)
						lowerCAmelCase__  :       Optional[Any]				=  np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371]				)

						assert np.abs(image_slice.flatten() - expected_slice				).max() < 5e-2



		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''				)
						lowerCAmelCase__  :       Optional[int]				=  sag_pipe.to(__lowerCamelCase				)
						sag_pipe.set_progress_bar_config(disable=__lowerCamelCase				)

						lowerCAmelCase__  :       Optional[int]				=  '''.'''
						lowerCAmelCase__  :       int				=  torch.manual_seed(0				)
						lowerCAmelCase__  :       Optional[int]				=  sag_pipe(
						    [prompt]	,width=7_68	,height=5_12	,generator=__lowerCamelCase	,guidance_scale=7.5	,sag_scale=1.0	,num_inference_steps=20	,output_type='''np'''	,)

						lowerCAmelCase__  :       List[str]				=  output.images

						assert image.shape == (1, 5_12, 7_68, 3)






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

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


__snake_case       :	Union[str, Any]            =logging.get_logger(__name__)

__snake_case       :	Dict            ={'vocab_file': 'sentencepiece.model'}

__snake_case       :	Optional[Any]            ={
    'vocab_file': {
        'google/rembert': 'https://huggingface.co/google/rembert/resolve/main/sentencepiece.model',
    },
}

__snake_case       :	int            ={
    'google/rembert': 2_5_6,
}



class  lowerCamelCase__    (							lowerCamelCase__):




		'''simple docstring'''



		snake_case_			=VOCAB_FILES_NAMES
		snake_case_			=PRETRAINED_VOCAB_FILES_MAP
		snake_case_			=PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES



		def __init__(self	,__lowerCamelCase	,__lowerCamelCase=False	,__lowerCamelCase=True	,__lowerCamelCase=True	,__lowerCamelCase="[CLS]"	,__lowerCamelCase="[SEP]"	,__lowerCamelCase="[UNK]"	,__lowerCamelCase="[SEP]"	,__lowerCamelCase="[PAD]"	,__lowerCamelCase="[CLS]"	,__lowerCamelCase="[MASK]"	,**__lowerCamelCase	,)  ->    Union[str, Any]:

						"""simple docstring"""






						super().__init__(
						    do_lower_case=__lowerCamelCase	,remove_space=__lowerCamelCase	,keep_accents=__lowerCamelCase	,bos_token=__lowerCamelCase	,eos_token=__lowerCamelCase	,unk_token=__lowerCamelCase	,sep_token=__lowerCamelCase	,pad_token=__lowerCamelCase	,cls_token=__lowerCamelCase	,mask_token=__lowerCamelCase	,**__lowerCamelCase	,)

						lowerCAmelCase__  :       int				=  do_lower_case
						lowerCAmelCase__  :       Optional[Any]				=  remove_space
						lowerCAmelCase__  :       Any				=  keep_accents
						lowerCAmelCase__  :       Dict				=  vocab_file

						lowerCAmelCase__  :       List[str]				=  spm.SentencePieceProcessor()
						self.sp_model.Load(__lowerCamelCase				)



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

						"""simple docstring"""






						return len(self.sp_model				)



		def 	lowerCAmelCase__   (self				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  {self.convert_ids_to_tokens(__lowerCamelCase				): i for i in range(self.vocab_size				)}
						vocab.update(self.added_tokens_encoder				)
						return vocab



		def __getstate__(self				)  ->    Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  self.__dict__.copy()
						lowerCAmelCase__  :       Dict				=  None
						return state



		def __setstate__(self	,__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  d
						lowerCAmelCase__  :       List[str]				=  spm.SentencePieceProcessor()
						self.sp_model.Load(self.vocab_file				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase=False				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  self.sp_model.EncodeAsPieces(__lowerCamelCase				)
						return pieces



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    Dict:

						"""simple docstring"""






						return self.sp_model.PieceToId(__lowerCamelCase				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    List[str]:

						"""simple docstring"""






						return self.sp_model.IdToPiece(__lowerCamelCase				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  self.sp_model.decode_pieces(__lowerCamelCase				)
						return out_string



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase = None				)  ->    List[int]:

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase = None	,__lowerCamelCase = False				)  ->    List[int]:

						"""simple docstring"""






						if already_has_special_tokens:
										if token_ids_a is not None:
														raise ValueError(
														    '''You should not supply a second sequence if the provided sequence of '''
														    '''ids is already formatted with special tokens for the model.'''				)
										return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

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



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase = None				)  ->    List[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  [self.sep_token_id]
						lowerCAmelCase__  :       Union[str, Any]				=  [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				) * [0] + len(token_ids_a + sep				) * [1]



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase = None				)  ->    Tuple[str]:

						"""simple docstring"""






						if not os.path.isdir(__lowerCamelCase				):
										logger.error('''Vocabulary path ({}) should be a directory'''.format(__lowerCamelCase				)				)
										return
						lowerCAmelCase__  :       Union[str, Any]				=  os.path.join(
						    __lowerCamelCase	,(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']				)

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

						return (out_vocab_file,)

