



'''simple docstring'''
import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    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 							lowercase__			(							lowercase		,      lowercase		,      unittest.TestCase     ):
	lowercase__						       =    StableDiffusionParadigmsPipeline
	lowercase__						       =    TEXT_TO_IMAGE_PARAMS
	lowercase__						       =    TEXT_TO_IMAGE_BATCH_PARAMS
	lowercase__						       =    TEXT_TO_IMAGE_IMAGE_PARAMS
	lowercase__						       =    TEXT_TO_IMAGE_IMAGE_PARAMS





	def 			UpperCamelCase_						(				self      :       Tuple				):







							'''simple docstring'''
							torch.manual_seed(0				)
							_UpperCamelCase      :	Optional[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				,attention_head_dim=(2, 4)				,use_linear_projection=lowerCamelCase__				,)
							_UpperCamelCase      :	str				       =     DDIMParallelScheduler(
							    beta_start=0.0_0_0_8_5				,beta_end=0.0_1_2				,beta_schedule='scaled_linear'				,clip_sample=lowerCamelCase__				,set_alpha_to_one=lowerCamelCase__				,)
							torch.manual_seed(0				)
							_UpperCamelCase      :	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				,sample_size=128				,)
							torch.manual_seed(0				)
							_UpperCamelCase      :	str				       =     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=1000				,hidden_act='gelu'				,projection_dim=512				,)
							_UpperCamelCase      :	Tuple				       =     CLIPTextModel(lowerCamelCase__				)
							_UpperCamelCase      :	int				       =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)

							_UpperCamelCase      :	Union[str, Any]				       =     {
							    'unet': unet,
							    'scheduler': scheduler,
							    'vae': vae,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'safety_checker': None,
							    'feature_extractor': None,
							}
							return components





	def 			UpperCamelCase_						(				self      :       int				,lowerCamelCase__      :       Union[str, Any]				,lowerCamelCase__      :       int=0				):







							'''simple docstring'''
							if str(lowerCamelCase__				).startswith('mps'				):
													_UpperCamelCase      :	Optional[int]				       =     torch.manual_seed(lowerCamelCase__				)
							else:
													_UpperCamelCase      :	List[str]				       =     torch.Generator(device=lowerCamelCase__				).manual_seed(lowerCamelCase__				)
							_UpperCamelCase      :	Optional[int]				       =     {
							    'prompt': 'a photograph of an astronaut riding a horse',
							    'generator': generator,
							    'num_inference_steps': 10,
							    'guidance_scale': 6.0,
							    'output_type': 'numpy',
							    'parallel': 3,
							    'debug': True,
							}
							return inputs





	def 			UpperCamelCase_						(				self      :       List[str]				):







							'''simple docstring'''
							_UpperCamelCase      :	int				       =     'cpu'  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase      :	List[str]				       =     self.get_dummy_components()
							_UpperCamelCase      :	Dict				       =     StableDiffusionParadigmsPipeline(**lowerCamelCase__				)
							_UpperCamelCase      :	int				       =     sd_pipe.to(lowerCamelCase__				)
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__				)

							_UpperCamelCase      :	Tuple				       =     self.get_dummy_inputs(lowerCamelCase__				)
							_UpperCamelCase      :	Optional[int]				       =     sd_pipe(**lowerCamelCase__				).images
							_UpperCamelCase      :	List[Any]				       =     image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							_UpperCamelCase      :	List[str]				       =     np.array([0.4_7_7_3, 0.5_4_1_7, 0.4_7_2_3, 0.4_9_2_5, 0.5_6_3_1, 0.4_7_5_2, 0.5_2_4_0, 0.4_9_3_5, 0.5_0_2_3]				)

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





	def 			UpperCamelCase_						(				self      :       int				):







							'''simple docstring'''
							_UpperCamelCase      :	Union[str, Any]				       =     'cpu'  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase      :	Optional[Any]				       =     self.get_dummy_components()
							torch.manual_seed(0				)
							_UpperCamelCase      :	Optional[int]				       =     DDPMParallelScheduler()
							torch.manual_seed(0				)
							_UpperCamelCase      :	Optional[Any]				       =     StableDiffusionParadigmsPipeline(**lowerCamelCase__				)
							_UpperCamelCase      :	Tuple				       =     sd_pipe.to(lowerCamelCase__				)
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__				)

							_UpperCamelCase      :	Optional[int]				       =     self.get_dummy_inputs(lowerCamelCase__				)
							_UpperCamelCase      :	Tuple				       =     sd_pipe(**lowerCamelCase__				).images
							_UpperCamelCase      :	Dict				       =     image[0, -3:, -3:, -1]
							assert image.shape == (1, 64, 64, 3)

							_UpperCamelCase      :	Optional[int]				       =     np.array([0.3_5_7_3, 0.4_4_2_0, 0.4_9_6_0, 0.4_7_9_9, 0.3_7_9_6, 0.3_8_7_9, 0.4_8_1_9, 0.4_3_6_5, 0.4_4_6_8]				)

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





	def 			UpperCamelCase_						(				self      :       Optional[int]				):







							'''simple docstring'''
							super().test_inference_batch_consistent(batch_sizes=[1, 2]				)





	def 			UpperCamelCase_						(				self      :       Optional[Any]				):







							'''simple docstring'''
							super().test_inference_batch_single_identical(batch_size=2				,expected_max_diff=3E-3				)





	def 			UpperCamelCase_						(				self      :       Any				):







							'''simple docstring'''
							_UpperCamelCase      :	int				       =     'cpu'  # ensure determinism for the device-dependent torch.Generator
							_UpperCamelCase      :	str				       =     self.get_dummy_components()
							_UpperCamelCase      :	int				       =     StableDiffusionParadigmsPipeline(**lowerCamelCase__				)
							_UpperCamelCase      :	Dict				       =     sd_pipe.to(lowerCamelCase__				)
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase__				)

							_UpperCamelCase      :	Any				       =     self.get_dummy_inputs(lowerCamelCase__				)
							_UpperCamelCase      :	str				       =     'french fries'
							_UpperCamelCase      :	List[Any]				       =     sd_pipe(**lowerCamelCase__				,negative_prompt=lowerCamelCase__				)
							_UpperCamelCase      :	str				       =     output.images
							_UpperCamelCase      :	Any				       =     image[0, -3:, -3:, -1]

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

							_UpperCamelCase      :	Union[str, Any]				       =     np.array([0.4_7_7_1, 0.5_4_2_0, 0.4_6_8_3, 0.4_9_1_8, 0.5_6_3_6, 0.4_7_2_5, 0.5_2_3_0, 0.4_9_2_3, 0.5_0_1_5]				)

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






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





	def 			UpperCamelCase_						(				self      :       Any				):







							'''simple docstring'''
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





	def 			UpperCamelCase_						(				self      :       str				,lowerCamelCase__      :       int=0				):







							'''simple docstring'''
							_UpperCamelCase      :	Tuple				       =     torch.Generator(device=lowerCamelCase__				).manual_seed(lowerCamelCase__				)
							_UpperCamelCase      :	str				       =     {
							    'prompt': 'a photograph of an astronaut riding a horse',
							    'generator': generator,
							    'num_inference_steps': 10,
							    'guidance_scale': 7.5,
							    'output_type': 'numpy',
							    'parallel': 3,
							    'debug': True,
							}
							return inputs





	def 			UpperCamelCase_						(				self      :       Optional[int]				):







							'''simple docstring'''
							_UpperCamelCase      :	Union[str, Any]				       =     'stabilityai/stable-diffusion-2-base'
							_UpperCamelCase      :	str				       =     DDIMParallelScheduler.from_pretrained(lowerCamelCase__				,subfolder='scheduler'				)
							_UpperCamelCase      :	Tuple				       =     StableDiffusionParadigmsPipeline.from_pretrained(lowerCamelCase__				,scheduler=lowerCamelCase__				,safety_checker=lowerCamelCase__				)
							pipe.to(lowerCamelCase__				)
							pipe.set_progress_bar_config(disable=lowerCamelCase__				)
							pipe.enable_attention_slicing()

							_UpperCamelCase      :	str				       =     self.get_inputs()
							_UpperCamelCase      :	List[Any]				       =     pipe(**lowerCamelCase__				).images
							_UpperCamelCase      :	Union[str, Any]				       =     image[0, -3:, -3:, -1].flatten()

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

							_UpperCamelCase      :	Any				       =     np.array([0.9_6_2_2, 0.9_6_0_2, 0.9_7_4_8, 0.9_5_9_1, 0.9_6_3_0, 0.9_6_9_1, 0.9_6_6_1, 0.9_6_3_1, 0.9_7_4_1]				)

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







'''simple docstring'''
import os
from shutil import copyfile
from typing import List, Optional, Tuple

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
					from .tokenization_mbartaa import MBartaaTokenizer
else:
					snake_case_				:			List[str]						 =						None


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

snake_case_				:			Dict						 =						{'vocab_file': 'sentencepiece.bpe.model', 'tokenizer_file': 'tokenizer.json'}

snake_case_				:			str						 =						{
    'vocab_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model'
        ),
    },
    'tokenizer_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/tokenizer.json'
        ),
    },
}

snake_case_				:			List[Any]						 =						{
    'facebook/mbart-large-50-one-to-many-mmt': 1024,
}

# fmt: off
snake_case_				:			Union[str, Any]						 =						['ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX', 'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP', 'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN', 'zh_CN', 'af_ZA', 'az_AZ', 'bn_IN', 'fa_IR', 'he_IL', 'hr_HR', 'id_ID', 'ka_GE', 'km_KH', 'mk_MK', 'ml_IN', 'mn_MN', 'mr_IN', 'pl_PL', 'ps_AF', 'pt_XX', 'sv_SE', 'sw_KE', 'ta_IN', 'te_IN', 'th_TH', 'tl_XX', 'uk_UA', 'ur_PK', 'xh_ZA', 'gl_ES', 'sl_SI']






class 							lowercase__			(							lowercase     ):
	lowercase__						       =    VOCAB_FILES_NAMES
	lowercase__						       =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	lowercase__						       =    PRETRAINED_VOCAB_FILES_MAP
	lowercase__						       =    ["""input_ids""", """attention_mask"""]
	lowercase__						       =    MBartaaTokenizer

	lowercase__						       =    []
	lowercase__						       =    []





	def __init__(				self      :       Optional[int]				,lowerCamelCase__      :       Union[str, Any]=None				,lowerCamelCase__      :       List[Any]=None				,lowerCamelCase__      :       Optional[Any]=None				,lowerCamelCase__      :       str=None				,lowerCamelCase__      :       List[str]="</s>"				,lowerCamelCase__      :       List[str]="</s>"				,lowerCamelCase__      :       Dict="<s>"				,lowerCamelCase__      :       int="<unk>"				,lowerCamelCase__      :       List[Any]="<pad>"				,lowerCamelCase__      :       Optional[int]="<mask>"				,**lowerCamelCase__      :       Optional[int]				,):







							'''simple docstring'''
							# Mask token behave like a normal word, i.e. include the space before it
							_UpperCamelCase      :	str				       =     AddedToken(lowerCamelCase__				,lstrip=lowerCamelCase__				,rstrip=lowerCamelCase__				) if isinstance(lowerCamelCase__				,lowerCamelCase__				) else mask_token

							_UpperCamelCase      :	Tuple				       =     kwargs.get('additional_special_tokens'				,[]				)
							kwargs["additional_special_tokens"] += [
							    code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
							]

							super().__init__(
							    lowerCamelCase__				,src_lang=lowerCamelCase__				,tgt_lang=lowerCamelCase__				,tokenizer_file=lowerCamelCase__				,eos_token=lowerCamelCase__				,sep_token=lowerCamelCase__				,cls_token=lowerCamelCase__				,unk_token=lowerCamelCase__				,pad_token=lowerCamelCase__				,mask_token=lowerCamelCase__				,**lowerCamelCase__				,)

							_UpperCamelCase      :	Optional[int]				       =     vocab_file
							_UpperCamelCase      :	Union[str, Any]				       =     False if not self.vocab_file else True

							_UpperCamelCase      :	Union[str, Any]				       =     {
							    lang_code: self.convert_tokens_to_ids(lowerCamelCase__				) for lang_code in FAIRSEQ_LANGUAGE_CODES
							}

							_UpperCamelCase      :	Union[str, Any]				       =     src_lang if src_lang is not None else 'en_XX'
							_UpperCamelCase      :	Any				       =     tgt_lang
							_UpperCamelCase      :	Optional[Any]				       =     self.lang_code_to_id[self._src_lang]
							self.set_src_lang_special_tokens(self._src_lang				)





	@property
	def 			UpperCamelCase_						(				self      :       Optional[Any]				):







							'''simple docstring'''
							return self._src_lang





	@src_lang.setter
	def 			UpperCamelCase_						(				self      :       Optional[Any]				,lowerCamelCase__      :       str				):







							'''simple docstring'''
							_UpperCamelCase      :	str				       =     new_src_lang
							self.set_src_lang_special_tokens(self._src_lang				)





	def 			UpperCamelCase_						(				self      :       Tuple				,lowerCamelCase__      :       List[int]				,lowerCamelCase__      :       Optional[List[int]] = None				):







							'''simple docstring'''
							if token_ids_a is None:
													return self.prefix_tokens + token_ids_a + self.suffix_tokens
							# We don't expect to process pairs, but leave the pair logic for API consistency
							return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens





	def 			UpperCamelCase_						(				self      :       Optional[Any]				,lowerCamelCase__      :       List[str]				,lowerCamelCase__      :       str = "en_XX"				,lowerCamelCase__      :       Optional[List[str]] = None				,lowerCamelCase__      :       str = "ro_RO"				,**lowerCamelCase__      :       Dict				,):







							'''simple docstring'''
							_UpperCamelCase      :	List[str]				       =     src_lang
							_UpperCamelCase      :	List[str]				       =     tgt_lang
							return super().prepare_seqaseq_batch(lowerCamelCase__				,lowerCamelCase__				,**lowerCamelCase__				)





	def 			UpperCamelCase_						(				self      :       Tuple				):







							'''simple docstring'''
							return self.set_src_lang_special_tokens(self.src_lang				)





	def 			UpperCamelCase_						(				self      :       Optional[int]				):







							'''simple docstring'''
							return self.set_tgt_lang_special_tokens(self.tgt_lang				)





	def 			UpperCamelCase_						(				self      :       Dict				,lowerCamelCase__      :       str				):







							'''simple docstring'''
							_UpperCamelCase      :	Any				       =     self.convert_tokens_to_ids(lowerCamelCase__				)
							_UpperCamelCase      :	List[Any]				       =     [self.cur_lang_code_id]
							_UpperCamelCase      :	List[Any]				       =     [self.eos_token_id]

							_UpperCamelCase      :	Optional[int]				       =     self.convert_ids_to_tokens(self.prefix_tokens				)
							_UpperCamelCase      :	List[Any]				       =     self.convert_ids_to_tokens(self.suffix_tokens				)

							_UpperCamelCase      :	Union[str, Any]				       =     processors.TemplateProcessing(
							    single=prefix_tokens_str + ['$A'] + suffix_tokens_str				,pair=prefix_tokens_str + ['$A', '$B'] + suffix_tokens_str				,special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str				,self.prefix_tokens + self.suffix_tokens				)				)				,)





	def 			UpperCamelCase_						(				self      :       Union[str, Any]				,lowerCamelCase__      :       str				):







							'''simple docstring'''
							_UpperCamelCase      :	List[str]				       =     self.convert_tokens_to_ids(lowerCamelCase__				)
							_UpperCamelCase      :	int				       =     [self.cur_lang_code_id]
							_UpperCamelCase      :	int				       =     [self.eos_token_id]

							_UpperCamelCase      :	Union[str, Any]				       =     self.convert_ids_to_tokens(self.prefix_tokens				)
							_UpperCamelCase      :	int				       =     self.convert_ids_to_tokens(self.suffix_tokens				)

							_UpperCamelCase      :	Dict				       =     processors.TemplateProcessing(
							    single=prefix_tokens_str + ['$A'] + suffix_tokens_str				,pair=prefix_tokens_str + ['$A', '$B'] + suffix_tokens_str				,special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str				,self.prefix_tokens + self.suffix_tokens				)				)				,)





	def 			UpperCamelCase_						(				self      :       Any				,lowerCamelCase__      :       Optional[int]				,lowerCamelCase__      :       str				,lowerCamelCase__      :       Optional[str]				,lowerCamelCase__      :       Optional[str]				,**lowerCamelCase__      :       List[Any]				):







							'''simple docstring'''
							if src_lang is None or tgt_lang is None:
													raise ValueError('Translation requires a `src_lang` and a `tgt_lang` for this model'				)
							_UpperCamelCase      :	Optional[int]				       =     src_lang
							_UpperCamelCase      :	Optional[Any]				       =     self(lowerCamelCase__				,add_special_tokens=lowerCamelCase__				,return_tensors=lowerCamelCase__				,**lowerCamelCase__				)
							_UpperCamelCase      :	Tuple				       =     self.convert_tokens_to_ids(lowerCamelCase__				)
							_UpperCamelCase      :	Optional[int]				       =     tgt_lang_id
							return inputs





	def 			UpperCamelCase_						(				self      :       Optional[int]				,lowerCamelCase__      :       str				,lowerCamelCase__      :       Optional[str] = None				):







							'''simple docstring'''
							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(lowerCamelCase__				):
													logger.error(F'Vocabulary path ({save_directory}) should be a directory'				)
													return
							_UpperCamelCase      :	Optional[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,)


