



"""simple docstring"""
import json
import os
from typing import Optional, Tuple

import regex as re

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


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

__SCREAMING_SNAKE_CASE  :						Union[str, Any]    							=      {
    "vocab_file": "vocab.json",
    "merges_file": "merges.txt",
}

__SCREAMING_SNAKE_CASE  :						Dict    							=      {
    "vocab_file": {"ctrl": "https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-vocab.json"},
    "merges_file": {"ctrl": "https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-merges.txt"},
}

__SCREAMING_SNAKE_CASE  :						Any    							=      {
    "ctrl": 2_5_6,
}

__SCREAMING_SNAKE_CASE  :						int    							=      {
    "Pregnancy": 1_6_8_6_2_9,
    "Christianity": 7_6_7_5,
    "Explain": 1_0_6_4_2_3,
    "Fitness": 6_3_4_4_0,
    "Saving": 6_3_1_6_3,
    "Ask": 2_7_1_7_1,
    "Ass": 9_5_9_8_5,
    "Joke": 1_6_3_5_0_9,
    "Questions": 4_5_6_2_2,
    "Thoughts": 4_9_6_0_5,
    "Retail": 5_2_3_4_2,
    "Feminism": 1_6_4_3_3_8,
    "Writing": 1_1_9_9_2,
    "Atheism": 1_9_2_2_6_3,
    "Netflix": 4_8_6_1_6,
    "Computing": 3_9_6_3_9,
    "Opinion": 4_3_2_1_3,
    "Alone": 4_4_9_6_7,
    "Funny": 5_8_9_1_7,
    "Gaming": 4_0_3_5_8,
    "Human": 4_0_8_8,
    "India": 1_3_3_1,
    "Joker": 7_7_1_3_8,
    "Diet": 3_6_2_0_6,
    "Legal": 1_1_8_5_9,
    "Norman": 4_9_3_9,
    "Tip": 7_2_6_8_9,
    "Weight": 5_2_3_4_3,
    "Movies": 4_6_2_7_3,
    "Running": 2_3_4_2_5,
    "Science": 2_0_9_0,
    "Horror": 3_7_7_9_3,
    "Confession": 6_0_5_7_2,
    "Finance": 1_2_2_5_0,
    "Politics": 1_6_3_6_0,
    "Scary": 1_9_1_9_8_5,
    "Support": 1_2_6_5_4,
    "Technologies": 3_2_5_1_6,
    "Teenage": 6_6_1_6_0,
    "Event": 3_2_7_6_9,
    "Learned": 6_7_4_6_0,
    "Notion": 1_8_2_7_7_0,
    "Wikipedia": 3_7_5_8_3,
    "Books": 6_6_6_5,
    "Extract": 7_6_0_5_0,
    "Confessions": 1_0_2_7_0_1,
    "Conspiracy": 7_5_9_3_2,
    "Links": 6_3_6_7_4,
    "Narcissus": 1_5_0_4_2_5,
    "Relationship": 5_4_7_6_6,
    "Relationships": 1_3_4_7_9_6,
    "Reviews": 4_1_6_7_1,
    "News": 4_2_5_6,
    "Translation": 2_6_8_2_0,
    "multilingual": 1_2_8_4_0_6,
}



def        lowerCAmelCase_(   lowercase_     :							Tuple							)	->       List[Any]:

				'''simple docstring'''





				_lowerCamelCase        =				set()
				_lowerCamelCase        =				word[0]
				for char in word[1:]:
								pairs.add((prev_char, char)							)
								_lowerCamelCase        =				char

				_lowerCamelCase        =				set(lowercase_							)
				return pairs




class 							lowerCamelCase_(   UpperCAmelCase_    ):





					'''simple docstring'''



					lowercase__   :       List[Any]										= VOCAB_FILES_NAMES
					lowercase__   :       Tuple										= PRETRAINED_VOCAB_FILES_MAP
					lowercase__   :       Optional[Any]										= PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					lowercase__   :       Union[str, Any]										= CONTROL_CODES

					def __init__(    self				,       lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__="<unk>"				,       **lowerCamelCase__   ):
									super().__init__(unk_token=_lowercase				,       **_lowercase   )

									with open(_lowercase				,       encoding='''utf-8'''   ) as vocab_handle:
													_lowerCamelCase        =				json.load(_lowercase   )
									_lowerCamelCase        =				{v: k for k, v in self.encoder.items()}
									with open(_lowercase				,       encoding='''utf-8'''   ) as merges_handle:
													_lowerCamelCase        =				merges_handle.read().split('''\n'''   )[1:-1]
									_lowerCamelCase        =				[tuple(merge.split()   ) for merge in merges]
									_lowerCamelCase        =				dict(zip(_lowercase				,       range(len(_lowercase   )   )   )   )
									_lowerCamelCase        =				{}

					@property
					def 			snake_case__      (    self   ):
									return len(self.encoder   )

					def 			snake_case__      (    self   ):
									return dict(self.encoder				,       **self.added_tokens_encoder   )

					def 			snake_case__      (    self				,       lowerCamelCase__   ):
									if token in self.cache:
													return self.cache[token]
									_lowerCamelCase        =				tuple(_lowercase   )
									_lowerCamelCase        =				tuple(list(word[:-1]   ) + [word[-1] + '''</w>''']   )
									_lowerCamelCase        =				get_pairs(_lowercase   )

									if not pairs:
													return token

									while True:
													_lowerCamelCase        =				min(_lowercase				,       key=lambda lowerCamelCase__   : self.bpe_ranks.get(_lowercase				,       float('''inf'''   )   )   )
													if bigram not in self.bpe_ranks:
																	break
													_lowerCamelCase        =				bigram
													_lowerCamelCase        =				[]
													_lowerCamelCase        =				0
													while i < len(_lowercase   ):
																	try:
																					_lowerCamelCase        =				word.index(_lowercase				,       _lowercase   )
																	except ValueError:
																					new_word.extend(word[i:]   )
																					break
																	else:
																					new_word.extend(word[i:j]   )
																					_lowerCamelCase        =				j

																	if word[i] == first and i < len(_lowercase   ) - 1 and word[i + 1] == second:
																					new_word.append(first + second   )
																					i += 2
																	else:
																					new_word.append(word[i]   )
																					i += 1
													_lowerCamelCase        =				tuple(_lowercase   )
													_lowerCamelCase        =				new_word
													if len(_lowercase   ) == 1:
																	break
													else:
																	_lowerCamelCase        =				get_pairs(_lowercase   )
									_lowerCamelCase        =				'@@ '.join(_lowercase   )
									_lowerCamelCase        =				word[:-4]
									_lowerCamelCase        =				word
									return word

					def 			snake_case__      (    self				,       lowerCamelCase__   ):
									_lowerCamelCase        =				[]

									_lowerCamelCase        =				re.findall(R'''\S+\n?'''				,       _lowercase   )

									for token in words:
													split_tokens.extend(list(self.bpe(_lowercase   ).split(''' '''   )   )   )
									return split_tokens

					def 			snake_case__      (    self				,       lowerCamelCase__   ):
									return self.encoder.get(_lowercase				,       self.encoder.get(self.unk_token   )   )

					def 			snake_case__      (    self				,       lowerCamelCase__   ):
									return self.decoder.get(_lowercase				,       self.unk_token   )

					def 			snake_case__      (    self				,       lowerCamelCase__   ):
									_lowerCamelCase        =				' '.join(_lowercase   ).replace('''@@ '''				,       ''''''   ).strip()
									return out_string





					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__ = 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''']   )
									_lowerCamelCase        =				os.path.join(
									    _lowercase				,       (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']   )

									with open(_lowercase				,       '''w'''				,       encoding='''utf-8'''   ) as f:
													f.write(json.dumps(self.encoder				,       indent=2				,       sort_keys=_lowercase				,       ensure_ascii=_lowercase   ) + '''\n'''   )

									_lowerCamelCase        =				0
									with open(_lowercase				,       '''w'''				,       encoding='''utf-8'''   ) as writer:
													writer.write('''#version: 0.2\n'''   )
													for bpe_tokens, token_index in sorted(self.bpe_ranks.items()				,       key=lambda lowerCamelCase__   : kv[1]   ):
																	if index != token_index:
																					logger.warning(
																					    F"""Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."""
																					    ''' Please check that the tokenizer is not corrupted!'''   )
																					_lowerCamelCase        =				token_index
																	writer.write(''' '''.join(_lowercase   ) + '''\n'''   )
																	index += 1

									return vocab_file, merge_file

					# def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True):
					#     filtered_tokens = ' '.join(self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens))
					#     tokens_generated_so_far = re.sub('(@@ )', '', string=filtered_tokens)
					#     tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)
					#     return ''.join(tokens_generated_so_far)





"""simple docstring"""
import gc
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class 							lowerCamelCase_(   A__,				A__,				unittest.TestCase    ):





					'''simple docstring'''



					lowercase__   :       Tuple										= StableDiffusionXLImgaImgPipeline
					lowercase__   :       List[str]										= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
					lowercase__   :       int										= PipelineTesterMixin.required_optional_params - {'latents'}
					lowercase__   :       Union[str, Any]										= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
					lowercase__   :       Dict										= IMAGE_TO_IMAGE_IMAGE_PARAMS
					lowercase__   :       List[str]										= IMAGE_TO_IMAGE_IMAGE_PARAMS

					def 			snake_case__      (    self   ):
									torch.manual_seed(0   )
									_lowerCamelCase        =				UNetaDConditionModel(
									    block_out_channels=(3_2, 6_4)				,       layers_per_block=2				,       sample_size=3_2				,       in_channels=4				,       out_channels=4				,       down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,       up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,       attention_head_dim=(2, 4)				,       use_linear_projection=lowerCamelCase__				,       addition_embed_type='''text_time'''				,       addition_time_embed_dim=8				,       transformer_layers_per_block=(1, 2)				,       projection_class_embeddings_input_dim=8_0				,       cross_attention_dim=6_4				,       )
									_lowerCamelCase        =				EulerDiscreteScheduler(
									    beta_start=0.0_0_0_8_5				,       beta_end=0.0_1_2				,       steps_offset=1				,       beta_schedule='''scaled_linear'''				,       timestep_spacing='''leading'''				,       )
									torch.manual_seed(0   )
									_lowerCamelCase        =				AutoencoderKL(
									    block_out_channels=[3_2, 6_4]				,       in_channels=3				,       out_channels=3				,       down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']				,       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']				,       latent_channels=4				,       sample_size=1_2_8				,       )
									torch.manual_seed(0   )
									_lowerCamelCase        =				CLIPTextConfig(
									    bos_token_id=0				,       eos_token_id=2				,       hidden_size=3_2				,       intermediate_size=3_7				,       layer_norm_eps=1e-05				,       num_attention_heads=4				,       num_hidden_layers=5				,       pad_token_id=1				,       vocab_size=1_0_0_0				,       hidden_act='''gelu'''				,       projection_dim=3_2				,       )
									_lowerCamelCase        =				CLIPTextModel(lowerCamelCase__   )
									_lowerCamelCase        =				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				,       local_files_only=lowerCamelCase__   )

									_lowerCamelCase        =				CLIPTextModelWithProjection(lowerCamelCase__   )
									_lowerCamelCase        =				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				,       local_files_only=lowerCamelCase__   )

									_lowerCamelCase        =				{
									    '''unet''': unet,
									    '''scheduler''': scheduler,
									    '''vae''': vae,
									    '''text_encoder''': text_encoder,
									    '''tokenizer''': tokenizer,
									    '''text_encoder_2''': text_encoder_a,
									    '''tokenizer_2''': tokenizer_a,
									    # "safety_checker": None,
									    # "feature_extractor": None,
									}
									return components

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__=0   ):
									_lowerCamelCase        =				floats_tensor((1, 3, 3_2, 3_2)				,       rng=random.Random(lowerCamelCase__   )   ).to(lowerCamelCase__   )
									_lowerCamelCase        =				image / 2 + 0.5
									if str(lowerCamelCase__   ).startswith('''mps'''   ):
													_lowerCamelCase        =				torch.manual_seed(lowerCamelCase__   )
									else:
													_lowerCamelCase        =				torch.Generator(device=lowerCamelCase__   ).manual_seed(lowerCamelCase__   )
									_lowerCamelCase        =				{
									    '''prompt''': '''A painting of a squirrel eating a burger''',
									    '''image''': image,
									    '''generator''': generator,
									    '''num_inference_steps''': 2,
									    '''guidance_scale''': 5.0,
									    '''output_type''': '''numpy''',
									    '''strength''': 0.7_5,
									}
									return inputs

					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				StableDiffusionXLImgaImgPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				sd_pipe.to(lowerCamelCase__   )
									sd_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				sd_pipe(**lowerCamelCase__   ).images
									_lowerCamelCase        =				image[0, -3:, -3:, -1]

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

									_lowerCamelCase        =				np.array([0.4_6_5_6, 0.4_8_4_0, 0.4_4_3_9, 0.6_6_9_8, 0.5_5_7_4, 0.4_5_2_4, 0.5_7_9_9, 0.5_9_4_3, 0.5_1_6_5]   )

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

					def 			snake_case__      (    self   ):
									super().test_attention_slicing_forward_pass(expected_max_diff=3e-3   )

					def 			snake_case__      (    self   ):
									super().test_inference_batch_single_identical(expected_max_diff=3e-3   )

					def 			snake_case__      (    self   ):
									pass





					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				self.get_dummy_components()
									_lowerCamelCase        =				StableDiffusionXLImgaImgPipeline(**lowerCamelCase__   )
									_lowerCamelCase        =				sd_pipe.to(lowerCamelCase__   )
									_lowerCamelCase        =				sd_pipe.to(lowerCamelCase__   )
									sd_pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									# forward without prompt embeds
									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				3 * ['''this is a negative prompt''']
									_lowerCamelCase        =				negative_prompt
									_lowerCamelCase        =				3 * [inputs['''prompt''']]

									_lowerCamelCase        =				sd_pipe(**lowerCamelCase__   )
									_lowerCamelCase        =				output.images[0, -3:, -3:, -1]

									# forward with prompt embeds
									_lowerCamelCase        =				self.get_dummy_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				3 * ['''this is a negative prompt''']
									_lowerCamelCase        =				3 * [inputs.pop('''prompt'''   )]

									(
									    (
									    _lowerCamelCase
									)  ,     (
									    _lowerCamelCase
									)  ,     (
									    _lowerCamelCase
									)  ,     (
									    _lowerCamelCase
									)  ,     
									)        =				sd_pipe.encode_prompt(lowerCamelCase__				,       negative_prompt=lowerCamelCase__   )

									_lowerCamelCase        =				sd_pipe(
									    **lowerCamelCase__				,       prompt_embeds=lowerCamelCase__				,       negative_prompt_embeds=lowerCamelCase__				,       pooled_prompt_embeds=lowerCamelCase__				,       negative_pooled_prompt_embeds=lowerCamelCase__				,       )
									_lowerCamelCase        =				output.images[0, -3:, -3:, -1]

									# make sure that it's equal
									assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()   ).max() < 1e-4


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





					'''simple docstring'''



					def 			snake_case__      (    self   ):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()

					def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__="cpu"				,       lowerCamelCase__=torch.floataa				,       lowerCamelCase__=0   ):
									_lowerCamelCase        =				torch.Generator(device=lowerCamelCase__   ).manual_seed(lowerCamelCase__   )
									_lowerCamelCase        =				np.random.RandomState(lowerCamelCase__   ).standard_normal((1, 4, 6_4, 6_4)   )
									_lowerCamelCase        =				torch.from_numpy(lowerCamelCase__   ).to(device=lowerCamelCase__				,       dtype=lowerCamelCase__   )
									_lowerCamelCase        =				{
									    '''prompt''': '''a photograph of an astronaut riding a horse''',
									    '''latents''': latents,
									    '''generator''': generator,
									    '''num_inference_steps''': 3,
									    '''guidance_scale''': 7.5,
									    '''output_type''': '''numpy''',
									}
									return inputs





					def 			snake_case__      (    self   ):
									_lowerCamelCase        =				DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''   )
									pipe.to(lowerCamelCase__   )
									pipe.set_progress_bar_config(disable=lowerCamelCase__   )

									_lowerCamelCase        =				self.get_inputs(lowerCamelCase__   )
									_lowerCamelCase        =				pipe(**lowerCamelCase__   ).images
									_lowerCamelCase        =				image[0, -3:, -3:, -1].flatten()

									assert image.shape == (1, 5_1_2, 5_1_2, 3)
									_lowerCamelCase        =				np.array([0.4_9_4_9_3, 0.4_7_8_9_6, 0.4_0_7_9_8, 0.5_4_2_1_4, 0.5_3_2_1_2, 0.4_8_2_0_2, 0.4_7_6_5_6, 0.4_6_3_2_9, 0.4_8_5_0_6]   )
									assert np.abs(image_slice - expected_slice   ).max() < 7e-3
