





'''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 __snake_case(				_lowerCAmelCase		,							_lowerCAmelCase		,							unittest.TestCase ):


						'''simple docstring'''



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


						def     __snake_case   ( self				)     ->     Union[str, Any]:
							torch.manual_seed(0				)
							lowerCAmelCase   =     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""")     ,	attention_head_dim=(2, 4)     ,	use_linear_projection=A_     ,	addition_embed_type="""text_time"""     ,	addition_time_embed_dim=8     ,	transformer_layers_per_block=(1, 2)     ,	projection_class_embeddings_input_dim=80     ,	cross_attention_dim=64     ,	)
							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=[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				)
							lowerCAmelCase   =     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=32     ,	)
							lowerCAmelCase   =     CLIPTextModel(A_				)
							lowerCAmelCase   =     CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     ,	local_files_only=A_				)

							lowerCAmelCase   =     CLIPTextModelWithProjection(A_				)
							lowerCAmelCase   =     CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     ,	local_files_only=A_				)

							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     ,	A_     ,	A_=0				)     ->     List[str]:
							lowerCAmelCase   =     floats_tensor((1, 3, 32, 32)     ,	rng=random.Random(A_				)				).to(A_				)
							lowerCAmelCase   =     image / 2 + 0.5
							if str(A_				).startswith("""mps"""				):
								lowerCAmelCase   =     torch.manual_seed(A_				)
							else:
								lowerCAmelCase   =     torch.Generator(device=A_				).manual_seed(A_				)
							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				)     ->     str:
							lowerCAmelCase   =     """cpu"""  # ensure determinism for the device-dependent torch.Generator
							lowerCAmelCase   =     self.get_dummy_components()
							lowerCAmelCase   =     StableDiffusionXLImgaImgPipeline(**A_				)
							lowerCAmelCase   =     sd_pipe.to(A_				)
							sd_pipe.set_progress_bar_config(disable=A_				)

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

							assert image.shape == (1, 32, 32, 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				)     ->     Union[str, Any]:
							super().test_attention_slicing_forward_pass(expected_max_diff=3e-3				)


						def     __snake_case   ( self				)     ->     Dict:
							super().test_inference_batch_single_identical(expected_max_diff=3e-3				)


						def     __snake_case   ( self				)     ->     Optional[int]:
							pass




						def     __snake_case   ( self				)     ->     List[str]:
							lowerCAmelCase   =     self.get_dummy_components()
							lowerCAmelCase   =     StableDiffusionXLImgaImgPipeline(**A_				)
							lowerCAmelCase   =     sd_pipe.to(A_				)
							lowerCAmelCase   =     sd_pipe.to(A_				)
							sd_pipe.set_progress_bar_config(disable=A_				)

							# forward without prompt embeds
							lowerCAmelCase   =     self.get_dummy_inputs(A_				)
							lowerCAmelCase   =     3 * ["""this is a negative prompt"""]
							lowerCAmelCase   =     negative_prompt
							lowerCAmelCase   =     3 * [inputs["""prompt"""]]

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

							# forward with prompt embeds
							lowerCAmelCase   =     self.get_dummy_inputs(A_				)
							lowerCAmelCase   =     3 * ["""this is a negative prompt"""]
							lowerCAmelCase   =     3 * [inputs.pop("""prompt"""				)]

							(
							    (
							    lowerCAmelCase       
							),       (
							    lowerCAmelCase       
							),       (
							    lowerCAmelCase       
							),       (
							    lowerCAmelCase       
							),              
							)   =     sd_pipe.encode_prompt(A_     ,	negative_prompt=A_				)

							lowerCAmelCase   =     sd_pipe(
							    **A_     ,	prompt_embeds=A_     ,	negative_prompt_embeds=A_     ,	pooled_prompt_embeds=A_     ,	negative_pooled_prompt_embeds=A_     ,	)
							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 __snake_case(				unittest.TestCase ):


						'''simple docstring'''



						def     __snake_case   ( self				)     ->     Optional[Any]:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()


						def     __snake_case   ( self     ,	A_     ,	A_="cpu"     ,	A_=torch.floataa     ,	A_=0				)     ->     Union[str, Any]:
							lowerCAmelCase   =     torch.Generator(device=A_				).manual_seed(A_				)
							lowerCAmelCase   =     np.random.RandomState(A_				).standard_normal((1, 4, 64, 64)				)
							lowerCAmelCase   =     torch.from_numpy(A_				).to(device=A_     ,	dtype=A_				)
							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				)     ->     Optional[int]:
							lowerCAmelCase   =     DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""				)
							pipe.to(A_				)
							pipe.set_progress_bar_config(disable=A_				)

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

							assert image.shape == (1, 512, 512, 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




'''simple docstring'''

import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import BatchEncoding, PreTrainedTokenizer
from ...utils import logging


UpperCAmelCase     		=   logging.get_logger(__name__)

UpperCAmelCase     		=   '▁'

UpperCAmelCase     		=   {
    'vocab_file': 'vocab.json',
    'spm_file': 'sentencepiece.bpe.model',
    'tokenizer_config_file': 'tokenizer_config.json',
}

UpperCAmelCase     		=   {
    'vocab_file': {
        'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json',
        'facebook/m2m100_1.2B': 'https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json',
    },
    'spm_file': {
        'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model',
        'facebook/m2m100_1.2B': 'https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model',
    },
    'tokenizer_config_file': {
        'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json',
        'facebook/m2m100_1.2B': 'https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json',
    },
}

UpperCAmelCase     		=   {
    'facebook/m2m100_418M': 1024,
}

# fmt: off
UpperCAmelCase     		=   {
    'm2m100': ['af', 'am', 'ar', 'ast', 'az', 'ba', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'ceb', 'cs', 'cy', 'da', 'de', 'el', 'en', 'es', 'et', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl', 'gu', 'ha', 'he', 'hi', 'hr', 'ht', 'hu', 'hy', 'id', 'ig', 'ilo', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'lb', 'lg', 'ln', 'lo', 'lt', 'lv', 'mg', 'mk', 'ml', 'mn', 'mr', 'ms', 'my', 'ne', 'nl', 'no', 'ns', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'ro', 'ru', 'sd', 'si', 'sk', 'sl', 'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'th', 'tl', 'tn', 'tr', 'uk', 'ur', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo', 'zh', 'zu'],
    'wmt21': ['en', 'ha', 'is', 'ja', 'cs', 'ru', 'zh', 'de']
}

class __snake_case(				_lowerCAmelCase ):


						'''simple docstring'''



						UpperCAmelCase							:   Optional[int]									=		VOCAB_FILES_NAMES
						UpperCAmelCase							:   Optional[Any]									=		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
						UpperCAmelCase							:   Optional[int]									=		PRETRAINED_VOCAB_FILES_MAP
						UpperCAmelCase							:   int									=		["input_ids", "attention_mask"]

						UpperCAmelCase							:   List[int]									=		[]
						UpperCAmelCase							:   List[int]									=		[]


						def __init__( self     ,	A_     ,	A_     ,	A_=None     ,	A_=None     ,	A_="<s>"     ,	A_="</s>"     ,	A_="</s>"     ,	A_="<pad>"     ,	A_="<unk>"     ,	A_="m2m100"     ,	A_ = None     ,	A_=8     ,	**A_     ,	)     ->     None:
							lowerCAmelCase   =     {} if sp_model_kwargs is None else sp_model_kwargs

							lowerCAmelCase   =     language_codes
							lowerCAmelCase   =     FAIRSEQ_LANGUAGE_CODES[language_codes]
							lowerCAmelCase   =     {lang_code: f'__{lang_code}__' for lang_code in fairseq_language_code}

							lowerCAmelCase   =     kwargs.get("""additional_special_tokens"""     ,	[]				)
							kwargs["additional_special_tokens"] += [
							    self.get_lang_token(A_				)
							    for lang_code in fairseq_language_code
							    if self.get_lang_token(A_				) not in kwargs["additional_special_tokens"]
							]

							super().__init__(
							    src_lang=A_     ,	tgt_lang=A_     ,	bos_token=A_     ,	eos_token=A_     ,	sep_token=A_     ,	unk_token=A_     ,	pad_token=A_     ,	language_codes=A_     ,	sp_model_kwargs=self.sp_model_kwargs     ,	num_madeup_words=A_     ,	**A_     ,	)

							lowerCAmelCase   =     vocab_file
							lowerCAmelCase   =     load_json(A_				)
							lowerCAmelCase   =     {v: k for k, v in self.encoder.items()}
							lowerCAmelCase   =     spm_file
							lowerCAmelCase   =     load_spm(A_     ,	self.sp_model_kwargs				)

							lowerCAmelCase   =     len(self.encoder				)

							lowerCAmelCase   =     {
							    self.get_lang_token(A_				): self.encoder_size + i for i, lang_code in enumerate(A_				)
							}
							lowerCAmelCase   =     {lang_code: self.encoder_size + i for i, lang_code in enumerate(A_				)}
							lowerCAmelCase   =     {v: k for k, v in self.lang_token_to_id.items()}

							lowerCAmelCase   =     src_lang if src_lang is not None else """en"""
							lowerCAmelCase   =     tgt_lang
							lowerCAmelCase   =     self.get_lang_id(self._src_lang				)
							self.set_src_lang_special_tokens(self._src_lang				)

							lowerCAmelCase   =     num_madeup_words


						@property
						def     __snake_case   ( self				)     ->     int:
							return len(self.encoder				) + len(self.lang_token_to_id				)


						@property
						def     __snake_case   ( self				)     ->     str:
							return self._src_lang


						@src_lang.setter
						def     __snake_case   ( self     ,	A_				)     ->     None:
							lowerCAmelCase   =     new_src_lang
							self.set_src_lang_special_tokens(self._src_lang				)


						def     __snake_case   ( self     ,	A_				)     ->     List[str]:
							return self.sp_model.encode(A_     ,	out_type=A_				)


						def     __snake_case   ( self     ,	A_				)     ->     Any:
							if token in self.lang_token_to_id:
								return self.lang_token_to_id[token]
							return self.encoder.get(A_     ,	self.encoder[self.unk_token]				)


						def     __snake_case   ( self     ,	A_				)     ->     str:
							if index in self.id_to_lang_token:
								return self.id_to_lang_token[index]
							return self.decoder.get(A_     ,	self.unk_token				)


						def     __snake_case   ( self     ,	A_				)     ->     List[str]:
							lowerCAmelCase   =     []
							lowerCAmelCase   =     """"""
							for token in tokens:
								# make sure that special tokens are not decoded using sentencepiece model
								if token in self.all_special_tokens:
									out_string += self.sp_model.decode(A_				) + token
									lowerCAmelCase   =     []
								else:
									current_sub_tokens.append(A_				)
							out_string += self.sp_model.decode(A_				)
							return out_string.strip()


						def     __snake_case   ( self     ,	A_     ,	A_ = None     ,	A_ = False				)     ->     List[int]:

							if already_has_special_tokens:
								return super().get_special_tokens_mask(
								    token_ids_a=A_     ,	token_ids_a=A_     ,	already_has_special_tokens=A_				)

							lowerCAmelCase   =     [1] * len(self.prefix_tokens				)
							lowerCAmelCase   =     [1] * len(self.suffix_tokens				)
							if token_ids_a is None:
								return prefix_ones + ([0] * len(A_				)) + suffix_ones
							return prefix_ones + ([0] * len(A_				)) + ([0] * len(A_				)) + suffix_ones


						def     __snake_case   ( self     ,	A_     ,	A_ = None				)     ->     List[int]:
							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     __snake_case   ( self				)     ->     Dict:
							lowerCAmelCase   =     {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				)     ->     Dict:
							lowerCAmelCase   =     self.__dict__.copy()
							lowerCAmelCase   =     None
							return state


						def __setstate__( self     ,	A_				)     ->     None:
							lowerCAmelCase   =     d

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

							lowerCAmelCase   =     load_spm(self.spm_file     ,	self.sp_model_kwargs				)


						def     __snake_case   ( self     ,	A_     ,	A_ = None				)     ->     Tuple[str]:
							lowerCAmelCase   =     Path(A_				)
							if not save_dir.is_dir():
								raise OSError(f'{save_directory} should be a directory'				)
							lowerCAmelCase   =     save_dir / (
							    (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""vocab_file"""]
							)
							lowerCAmelCase   =     save_dir / (
							    (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""spm_file"""]
							)

							save_json(self.encoder     ,	A_				)

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

							return (str(A_				), str(A_				))


						def     __snake_case   ( self     ,	A_     ,	A_ = "en"     ,	A_ = None     ,	A_ = "ro"     ,	**A_     ,	)     ->     BatchEncoding:
							lowerCAmelCase   =     src_lang
							lowerCAmelCase   =     tgt_lang
							self.set_src_lang_special_tokens(self.src_lang				)
							return super().prepare_seqaseq_batch(A_     ,	A_     ,	**A_				)


						def     __snake_case   ( self     ,	A_     ,	A_     ,	A_     ,	**A_				)     ->     str:
							if src_lang is None or tgt_lang is None:
								raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model"""				)
							lowerCAmelCase   =     src_lang
							lowerCAmelCase   =     self(A_     ,	add_special_tokens=A_     ,	**A_				)
							lowerCAmelCase   =     self.get_lang_id(A_				)
							lowerCAmelCase   =     tgt_lang_id
							return inputs


						def     __snake_case   ( self				)     ->     Any:
							self.set_src_lang_special_tokens(self.src_lang				)


						def     __snake_case   ( self				)     ->     Optional[int]:
							self.set_tgt_lang_special_tokens(self.tgt_lang				)


						def     __snake_case   ( self     ,	A_				)     ->     None:
							lowerCAmelCase   =     self.get_lang_token(A_				)
							lowerCAmelCase   =     self.lang_token_to_id[lang_token]
							lowerCAmelCase   =     [self.cur_lang_id]
							lowerCAmelCase   =     [self.eos_token_id]


						def     __snake_case   ( self     ,	A_				)     ->     None:
							lowerCAmelCase   =     self.get_lang_token(A_				)
							lowerCAmelCase   =     self.lang_token_to_id[lang_token]
							lowerCAmelCase   =     [self.cur_lang_id]
							lowerCAmelCase   =     [self.eos_token_id]


						def     __snake_case   ( self     ,	A_				)     ->     str:
							return self.lang_code_to_token[lang]




						def     __snake_case   ( self     ,	A_				)     ->     int:
							lowerCAmelCase   =     self.get_lang_token(A_				)
							return self.lang_token_to_id[lang_token]







def 		_snake_case		( _SCREAMING_SNAKE_CASE					:						str							,		_SCREAMING_SNAKE_CASE					:						Dict[str, Any]							)      ->						sentencepiece.SentencePieceProcessor:







	"""simple docstring"""
	lowerCAmelCase   =     sentencepiece.SentencePieceProcessor(**_SCREAMING_SNAKE_CASE							)
	spm.Load(str(_SCREAMING_SNAKE_CASE							)							)
	return spm






def 		_snake_case		( _SCREAMING_SNAKE_CASE					:						str							)      ->						Union[Dict, List]:







	"""simple docstring"""
	with open(_SCREAMING_SNAKE_CASE							,		"""r"""							) as f:
		return json.load(_SCREAMING_SNAKE_CASE							)






def 		_snake_case		( _SCREAMING_SNAKE_CASE					:						Dict							,		_SCREAMING_SNAKE_CASE					:						str							)      ->						None:







	"""simple docstring"""
	with open(_SCREAMING_SNAKE_CASE							,		"""w"""							) as f:
		json.dump(_SCREAMING_SNAKE_CASE							,		_SCREAMING_SNAKE_CASE							,		indent=2							)