
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.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu





class 							UpperCAmelCase      (  unittest.TestCase  ):




	'''simple docstring'''





	def 				UpperCamelCase_  (       self					: int ):
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()






	@property
	def 				UpperCamelCase_  (       self					: str ):
					__A   			= 1
					__A   			= 3
					__A   			= (32, 32)

					__A   			= floats_tensor((batch_size, num_channels) + sizes							,rng=random.Random(0 ) ).to(A )
					return image






	@property
	def 				UpperCamelCase_  (       self					: Dict ):
					torch.manual_seed(0 )
					__A   			= 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 				UpperCamelCase_  (       self					: Any ):
					torch.manual_seed(0 )
					__A   			= 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 				UpperCamelCase_  (       self					: List[str] ):
					torch.manual_seed(0 )
					__A   			= 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							,)
					return CLIPTextModel(A )






	@property
	def 				UpperCamelCase_  (       self					: Tuple ):
					def extract(*A					: Any							,**A					: Optional[int] ):





									class 							UpperCAmelCase      :




										'''simple docstring'''





										def __init__(       self					: List[Any] ):
														__A   			= torch.ones([0] )







										def 				UpperCamelCase_  (       self					: List[Any]							,A					: Optional[int] ):
														self.pixel_values.to(A )
														return self





									return Out()

					return extract






	def 				UpperCamelCase_  (       self					: Tuple ):
					__A   			= "cpu"  # ensure determinism for the device-dependent torch.Generator
					__A   			= self.dummy_cond_unet
					__A   			= DDIMScheduler(
					    beta_start=0.0_00_85							,beta_end=0.0_12							,beta_schedule="scaled_linear"							,clip_sample=A							,set_alpha_to_one=A							,)

					__A   			= self.dummy_vae
					__A   			= self.dummy_text_encoder
					__A   			= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

					# make sure here that pndm scheduler skips prk
					__A   			= StableDiffusionPipeline(
					    unet=A							,scheduler=A							,vae=A							,text_encoder=A							,tokenizer=A							,safety_checker=A							,feature_extractor=self.dummy_extractor							,)
					__A   			= sd_pipe.to(A )
					sd_pipe.set_progress_bar_config(disable=A )

					__A   			= "A painting of a squirrel eating a burger"

					__A   			= torch.Generator(device=A ).manual_seed(0 )
					__A   			= sd_pipe([prompt]							,generator=A							,guidance_scale=6.0							,num_inference_steps=2							,output_type="np" )
					__A   			= output.images

					__A   			= torch.Generator(device=A ).manual_seed(0 )
					__A   			= sd_pipe(
					    [prompt]							,generator=A							,guidance_scale=6.0							,num_inference_steps=2							,output_type="np"							,return_dict=A							,)[0]

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

					assert image.shape == (1, 64, 64, 3)
					__A   			= np.array([0.57_56, 0.61_18, 0.50_05, 0.50_41, 0.54_71, 0.47_26, 0.49_76, 0.48_65, 0.48_64] )

					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 				UpperCamelCase_  (       self					: Any ):
					__A   			= "cpu"  # ensure determinism for the device-dependent torch.Generator
					__A   			= self.dummy_cond_unet
					__A   			= PNDMScheduler(skip_prk_steps=A )
					__A   			= self.dummy_vae
					__A   			= self.dummy_text_encoder
					__A   			= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

					# make sure here that pndm scheduler skips prk
					__A   			= StableDiffusionPipeline(
					    unet=A							,scheduler=A							,vae=A							,text_encoder=A							,tokenizer=A							,safety_checker=A							,feature_extractor=self.dummy_extractor							,)
					__A   			= sd_pipe.to(A )
					sd_pipe.set_progress_bar_config(disable=A )

					__A   			= "A painting of a squirrel eating a burger"
					__A   			= torch.Generator(device=A ).manual_seed(0 )
					__A   			= sd_pipe([prompt]							,generator=A							,guidance_scale=6.0							,num_inference_steps=2							,output_type="np" )

					__A   			= output.images

					__A   			= torch.Generator(device=A ).manual_seed(0 )
					__A   			= sd_pipe(
					    [prompt]							,generator=A							,guidance_scale=6.0							,num_inference_steps=2							,output_type="np"							,return_dict=A							,)[0]

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

					assert image.shape == (1, 64, 64, 3)
					__A   			= np.array([0.51_25, 0.57_16, 0.48_28, 0.50_60, 0.56_50, 0.47_68, 0.51_85, 0.48_95, 0.49_93] )

					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 				UpperCamelCase_  (       self					: Optional[int] ):
					__A   			= StableDiffusionPipeline.from_pretrained(
					    "hf-internal-testing/tiny-stable-diffusion-lms-pipe"							,safety_checker=A )
					assert isinstance(A							,A )
					assert isinstance(pipe.scheduler							,A )
					assert pipe.safety_checker is None

					__A   			= 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(A )
									__A   			= StableDiffusionPipeline.from_pretrained(A )

					# sanity check that the pipeline still works
					assert pipe.safety_checker is None
					__A   			= 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 				UpperCamelCase_  (       self					: Dict ):
					__A   			= self.dummy_cond_unet
					__A   			= PNDMScheduler(skip_prk_steps=A )
					__A   			= self.dummy_vae
					__A   			= self.dummy_text_encoder
					__A   			= CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

					# put models in fp16
					__A   			= unet.half()
					__A   			= vae.half()
					__A   			= bert.half()

					# make sure here that pndm scheduler skips prk
					__A   			= StableDiffusionPipeline(
					    unet=A							,scheduler=A							,vae=A							,text_encoder=A							,tokenizer=A							,safety_checker=A							,feature_extractor=self.dummy_extractor							,)
					__A   			= sd_pipe.to(A )
					sd_pipe.set_progress_bar_config(disable=A )

					__A   			= "A painting of a squirrel eating a burger"
					__A   			= sd_pipe([prompt]							,num_inference_steps=2							,output_type="np" ).images

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





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




	'''simple docstring'''





	def 				UpperCamelCase_  (       self					: List[str] ):
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()






	def 				UpperCamelCase_  (       self					: Dict ):
					__A   			= StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"							,safety_checker=A )
					__A   			= LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config )
					__A   			= sd_pipe.to(A )
					sd_pipe.set_progress_bar_config(disable=A )

					__A   			= (
					    "portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle"
					    " coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with"
					    " anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and"
					    " children from bahnhof zoo, detailed "
					)
					__A   			= 40_03_66_03_46
					__A   			= 7

					# without safety guidance (sld_guidance_scale = 0)
					__A   			= torch.manual_seed(A )
					__A   			= sd_pipe(
					    [prompt]							,generator=A							,guidance_scale=A							,num_inference_steps=50							,output_type="np"							,width=5_12							,height=5_12							,sld_guidance_scale=0							,)

					__A   			= output.images
					__A   			= image[0, -3:, -3:, -1]
					__A   			= [0.22_78, 0.22_31, 0.22_49, 0.23_33, 0.23_03, 0.18_85, 0.22_73, 0.21_44, 0.21_76]

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

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

					# without safety guidance (strong configuration)
					__A   			= torch.manual_seed(A )
					__A   			= sd_pipe(
					    [prompt]							,generator=A							,guidance_scale=A							,num_inference_steps=50							,output_type="np"							,width=5_12							,height=5_12							,sld_guidance_scale=20_00							,sld_warmup_steps=7							,sld_threshold=0.0_25							,sld_momentum_scale=0.5							,sld_mom_beta=0.7							,)

					__A   			= output.images
					__A   			= image[0, -3:, -3:, -1]
					__A   			= [0.23_83, 0.22_76, 0.2_36, 0.21_92, 0.21_86, 0.20_53, 0.19_71, 0.19_01, 0.17_19]

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

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






	def 				UpperCamelCase_  (       self					: List[str] ):
					__A   			= StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"							,safety_checker=A )
					__A   			= LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config )
					__A   			= sd_pipe.to(A )
					sd_pipe.set_progress_bar_config(disable=A )

					__A   			= "padme amidala taking a bath artwork, safe for work, no nudity"
					__A   			= 27_34_97_17_55
					__A   			= 7

					__A   			= torch.manual_seed(A )
					__A   			= sd_pipe(
					    [prompt]							,generator=A							,guidance_scale=A							,num_inference_steps=50							,output_type="np"							,width=5_12							,height=5_12							,sld_guidance_scale=0							,)

					__A   			= output.images
					__A   			= image[0, -3:, -3:, -1]
					__A   			= [0.35_02, 0.36_22, 0.33_96, 0.36_42, 0.34_78, 0.33_18, 0.35, 0.33_48, 0.32_97]

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

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

					__A   			= torch.manual_seed(A )
					__A   			= sd_pipe(
					    [prompt]							,generator=A							,guidance_scale=A							,num_inference_steps=50							,output_type="np"							,width=5_12							,height=5_12							,sld_guidance_scale=20_00							,sld_warmup_steps=7							,sld_threshold=0.0_25							,sld_momentum_scale=0.5							,sld_mom_beta=0.7							,)

					__A   			= output.images
					__A   			= image[0, -3:, -3:, -1]
					__A   			= [0.55_31, 0.52_06, 0.48_95, 0.51_56, 0.51_82, 0.47_51, 0.48_02, 0.48_03, 0.44_43]

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

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







	def 				UpperCamelCase_  (       self					: Dict ):
					__A   			= StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5" )
					__A   			= sd_pipe.to(A )
					sd_pipe.set_progress_bar_config(disable=A )

					__A   			= (
					    "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."
					    " leyendecker"
					)
					__A   			= 10_44_35_52_34
					__A   			= 12

					__A   			= torch.manual_seed(A )
					__A   			= sd_pipe(
					    [prompt]							,generator=A							,guidance_scale=A							,num_inference_steps=50							,output_type="np"							,width=5_12							,height=5_12							,sld_guidance_scale=0							,)

					__A   			= output.images
					__A   			= image[0, -3:, -3:, -1]
					__A   			= np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] )

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

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

					__A   			= torch.manual_seed(A )
					__A   			= sd_pipe(
					    [prompt]							,generator=A							,guidance_scale=A							,num_inference_steps=50							,output_type="np"							,width=5_12							,height=5_12							,sld_guidance_scale=20_00							,sld_warmup_steps=7							,sld_threshold=0.0_25							,sld_momentum_scale=0.5							,sld_mom_beta=0.7							,)

					__A   			= output.images
					__A   			= image[0, -3:, -3:, -1]
					__A   			= np.array([0.58_18, 0.62_85, 0.68_35, 0.60_19, 0.6_25, 0.67_54, 0.60_96, 0.63_34, 0.65_61] )
					assert image.shape == (1, 5_12, 5_12, 3)

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


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


SCREAMING_SNAKE_CASE							:List[Any]							 =				logging.get_logger(__name__)

SCREAMING_SNAKE_CASE							:int							 =				{'vocab_file': 'spiece.model'}

SCREAMING_SNAKE_CASE							:Union[str, Any]							 =				{
    'vocab_file': {
        'google/bigbird-roberta-base': 'https://huggingface.co/google/bigbird-roberta-base/resolve/main/spiece.model',
        'google/bigbird-roberta-large': (
            'https://huggingface.co/google/bigbird-roberta-large/resolve/main/spiece.model'
        ),
        'google/bigbird-base-trivia-itc': (
            'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/spiece.model'
        ),
    }
}

SCREAMING_SNAKE_CASE							:int							 =				{
    'google/bigbird-roberta-base': 4096,
    'google/bigbird-roberta-large': 4096,
    'google/bigbird-base-trivia-itc': 4096,
}





class 							UpperCAmelCase      (  __SCREAMING_SNAKE_CASE  ):




	'''simple docstring'''





	snake_case_					    =     VOCAB_FILES_NAMES
	snake_case_					    =     PRETRAINED_VOCAB_FILES_MAP
	snake_case_					    =     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	snake_case_					    =     ["input_ids", "attention_mask"]
	snake_case_					    =     []






	def __init__(       self					: Any							,A					: List[str]							,A					: str="<unk>"							,A					: int="<s>"							,A					: Union[str, Any]="</s>"							,A					: List[str]="<pad>"							,A					: int="[SEP]"							,A					: Optional[Any]="[MASK]"							,A					: Tuple="[CLS]"							,A					: Optional[Dict[str, Any]] = None							,**A					: Any							,):
					__A   			= AddedToken(A							,lstrip=A							,rstrip=A ) if isinstance(A							,A ) else bos_token
					__A   			= AddedToken(A							,lstrip=A							,rstrip=A ) if isinstance(A							,A ) else eos_token
					__A   			= AddedToken(A							,lstrip=A							,rstrip=A ) if isinstance(A							,A ) else unk_token
					__A   			= AddedToken(A							,lstrip=A							,rstrip=A ) if isinstance(A							,A ) else pad_token
					__A   			= AddedToken(A							,lstrip=A							,rstrip=A ) if isinstance(A							,A ) else cls_token
					__A   			= AddedToken(A							,lstrip=A							,rstrip=A ) if isinstance(A							,A ) else sep_token

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

					__A   			= {} if sp_model_kwargs is None else sp_model_kwargs

					super().__init__(
					    bos_token=A							,eos_token=A							,unk_token=A							,pad_token=A							,sep_token=A							,mask_token=A							,cls_token=A							,sp_model_kwargs=self.sp_model_kwargs							,**A							,)

					__A   			= vocab_file

					__A   			= spm.SentencePieceProcessor(**self.sp_model_kwargs )
					self.sp_model.Load(A )






	@property
	def 				UpperCamelCase_  (       self					: List[str] ):
					return self.sp_model.get_piece_size()






	def 				UpperCamelCase_  (       self					: Optional[Any] ):
					__A   			= {self.convert_ids_to_tokens(A ): i for i in range(self.vocab_size )}
					vocab.update(self.added_tokens_encoder )
					return vocab






	def __getstate__(       self					: Optional[int] ):
					__A   			= self.__dict__.copy()
					__A   			= None
					return state






	def __setstate__(       self					: str							,A					: Optional[Any] ):
					__A   			= d

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

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






	def 				UpperCamelCase_  (       self					: Any							,A					: str ):
					return self.sp_model.encode(A							,out_type=A )






	def 				UpperCamelCase_  (       self					: List[str]							,A					: Tuple ):
					return self.sp_model.piece_to_id(A )






	def 				UpperCamelCase_  (       self					: List[Any]							,A					: Tuple ):
					__A   			= self.sp_model.IdToPiece(A )
					return token






	def 				UpperCamelCase_  (       self					: List[Any]							,A					: int ):
					__A   			= []
					__A   			= ""
					__A   			= 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(A ) + token
													__A   			= True
													__A   			= []
									else:
													current_sub_tokens.append(A )
													__A   			= False
					out_string += self.sp_model.decode(A )
					return out_string.strip()






	def 				UpperCamelCase_  (       self					: Tuple							,A					: List[int]							,A					: bool = False							,A					: bool = None							,A					: bool = True							,**A					: Union[str, Any]							,):
					__A   			= kwargs.pop("use_source_tokenizer"							,A )

					__A   			= self.convert_ids_to_tokens(A							,skip_special_tokens=A )

					# 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
					__A   			= []
					__A   			= []
					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(A ) )
																	__A   			= []
													sub_texts.append(A )
									else:
													current_sub_text.append(A )
					if current_sub_text:
									sub_texts.append(self.convert_tokens_to_string(A ) )

					# Mimic the behavior of the Rust tokenizer:
					# No space before [MASK] and [SEP]
					if spaces_between_special_tokens:
									__A   			= re.sub(R" (\[(MASK|SEP)\])"							,R"\1"							," ".join(A ) )
					else:
									__A   			= "".join(A )

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






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

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

					return (out_vocab_file,)






	def 				UpperCamelCase_  (       self					: Dict							,A					: List[int]							,A					: Optional[List[int]] = None ):
					if token_ids_a is None:
									return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
					__A   			= [self.cls_token_id]
					__A   			= [self.sep_token_id]
					return cls + token_ids_a + sep + token_ids_a + sep






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

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







	def 				UpperCamelCase_  (       self					: Any							,A					: List[int]							,A					: Optional[List[int]] = None ):
					__A   			= [self.sep_token_id]
					__A   			= [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]
