





import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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 PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class   a_ (			__a		,							__a		,							__a		,							unittest.TestCase      ):
			A__      :			Any       	=  AltDiffusionPipeline
			A__      :			int       	=  TEXT_TO_IMAGE_PARAMS
			A__      :			int       	=  TEXT_TO_IMAGE_BATCH_PARAMS
			A__      :			Union[str, Any]       	=  TEXT_TO_IMAGE_IMAGE_PARAMS
			A__      :			Optional[Any]       	=  TEXT_TO_IMAGE_IMAGE_PARAMS

			def      lowerCAmelCase(   self   :  Union[str, Any]						):






							"""simple docstring"""
							torch.manual_seed(0						)
							snake_case   :					Union[str, 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     , )
							snake_case   :					Tuple    						=				DDIMScheduler(
							    beta_start=0.0_0085     , beta_end=0.012     , beta_schedule='''scaled_linear'''     , clip_sample=__UpperCamelCase     , set_alpha_to_one=__UpperCamelCase     , )
							torch.manual_seed(0						)
							snake_case   :					Union[str, 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     , )

							# TODO: address the non-deterministic text encoder (fails for save-load tests)
							# torch.manual_seed(0)
							# text_encoder_config = RobertaSeriesConfig(
							#     hidden_size=32,
							#     project_dim=32,
							#     intermediate_size=37,
							#     layer_norm_eps=1e-05,
							#     num_attention_heads=4,
							#     num_hidden_layers=5,
							#     vocab_size=5002,
							# )
							# text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

							torch.manual_seed(0						)
							snake_case   :					List[str]    						=				CLIPTextConfig(
							    bos_token_id=0     , eos_token_id=2     , hidden_size=32     , projection_dim=32     , intermediate_size=37     , layer_norm_eps=1e-0_5     , num_attention_heads=4     , num_hidden_layers=5     , pad_token_id=1     , vocab_size=5_002     , )
							snake_case   :					Tuple    						=				CLIPTextModel(__UpperCamelCase						)

							snake_case   :					Optional[int]    						=				XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''						)
							snake_case   :					Union[str, Any]    						=				77

							snake_case   :					Any    						=				{
							    """unet""": unet,
							    """scheduler""": scheduler,
							    """vae""": vae,
							    """text_encoder""": text_encoder,
							    """tokenizer""": tokenizer,
							    """safety_checker""": None,
							    """feature_extractor""": None,
							}
							return components

			def      lowerCAmelCase(   self   :  Any     , UpperCAmelCase__   :  Any     , UpperCAmelCase__   :  List[Any]=0						):






							"""simple docstring"""
							if str(__UpperCamelCase						).startswith('''mps'''						):
											snake_case   :					Optional[Any]    						=				torch.manual_seed(__UpperCamelCase						)
							else:
											snake_case   :					Optional[Any]    						=				torch.Generator(device=__UpperCamelCase						).manual_seed(__UpperCamelCase						)
							snake_case   :					str    						=				{
							    """prompt""": """A painting of a squirrel eating a burger""",
							    """generator""": generator,
							    """num_inference_steps""": 2,
							    """guidance_scale""": 6.0,
							    """output_type""": """numpy""",
							}
							return inputs

			def      lowerCAmelCase(   self   :  int						):






							"""simple docstring"""
							super().test_attention_slicing_forward_pass(expected_max_diff=3e-3						)

			def      lowerCAmelCase(   self   :  List[Any]						):






							"""simple docstring"""
							super().test_inference_batch_single_identical(expected_max_diff=3e-3						)

			def      lowerCAmelCase(   self   :  Tuple						):






							"""simple docstring"""
							snake_case   :					str    						=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator

							snake_case   :					Dict    						=				self.get_dummy_components()
							torch.manual_seed(0						)
							snake_case   :					Any    						=				RobertaSeriesConfig(
							    hidden_size=32     , project_dim=32     , intermediate_size=37     , layer_norm_eps=1e-0_5     , num_attention_heads=4     , num_hidden_layers=5     , vocab_size=5_002     , )
							# TODO: remove after fixing the non-deterministic text encoder
							snake_case   :					Tuple    						=				RobertaSeriesModelWithTransformation(__UpperCamelCase						)
							snake_case   :					List[Any]    						=				text_encoder

							snake_case   :					Optional[int]    						=				AltDiffusionPipeline(**__UpperCamelCase						)
							snake_case   :					Union[str, Any]    						=				alt_pipe.to(__UpperCamelCase						)
							alt_pipe.set_progress_bar_config(disable=__UpperCamelCase						)

							snake_case   :					int    						=				self.get_dummy_inputs(__UpperCamelCase						)
							snake_case   :					List[str]    						=				"""A photo of an astronaut"""
							snake_case   :					Dict    						=				alt_pipe(**__UpperCamelCase						)
							snake_case   :					int    						=				output.images
							snake_case   :					str    						=				image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							snake_case   :					str    						=				np.array(
							    [0.574_8162, 0.6044_7145, 0.4882_1217, 0.5010_0636, 0.543_1185, 0.4576_3683, 0.4965_7696, 0.4813_2733, 0.4757_3093]						)

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



			def      lowerCAmelCase(   self   :  Optional[int]						):






							"""simple docstring"""
							snake_case   :					Tuple    						=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator

							snake_case   :					Dict    						=				self.get_dummy_components()
							snake_case   :					Dict    						=				PNDMScheduler(skip_prk_steps=__UpperCamelCase						)
							torch.manual_seed(0						)
							snake_case   :					Any    						=				RobertaSeriesConfig(
							    hidden_size=32     , project_dim=32     , intermediate_size=37     , layer_norm_eps=1e-0_5     , num_attention_heads=4     , num_hidden_layers=5     , vocab_size=5_002     , )
							# TODO: remove after fixing the non-deterministic text encoder
							snake_case   :					Any    						=				RobertaSeriesModelWithTransformation(__UpperCamelCase						)
							snake_case   :					Any    						=				text_encoder
							snake_case   :					str    						=				AltDiffusionPipeline(**__UpperCamelCase						)
							snake_case   :					Tuple    						=				alt_pipe.to(__UpperCamelCase						)
							alt_pipe.set_progress_bar_config(disable=__UpperCamelCase						)

							snake_case   :					Optional[Any]    						=				self.get_dummy_inputs(__UpperCamelCase						)
							snake_case   :					Any    						=				alt_pipe(**__UpperCamelCase						)
							snake_case   :					Dict    						=				output.images
							snake_case   :					Optional[int]    						=				image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							snake_case   :					int    						=				np.array(
							    [0.5160_5093, 0.570_7241, 0.4736_5507, 0.5057_8886, 0.563_3877, 0.464_2503, 0.518_2081, 0.4876_3484, 0.4908_4237]						)

							assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
@slow
@require_torch_gpu
class   a_ (			unittest.TestCase      ):

			def      lowerCAmelCase(   self   :  Union[str, Any]						):






							"""simple docstring"""
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()

			def      lowerCAmelCase(   self   :  Dict						):






							"""simple docstring"""
							# make sure here that pndm scheduler skips prk
							snake_case   :					Union[str, Any]    						=				AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''     , safety_checker=__UpperCamelCase						)
							snake_case   :					Optional[int]    						=				alt_pipe.to(__UpperCamelCase						)
							alt_pipe.set_progress_bar_config(disable=__UpperCamelCase						)

							snake_case   :					List[Any]    						=				"""A painting of a squirrel eating a burger"""
							snake_case   :					Optional[Any]    						=				torch.manual_seed(0						)
							snake_case   :					Union[str, Any]    						=				alt_pipe([prompt]     , generator=__UpperCamelCase     , guidance_scale=6.0     , num_inference_steps=20     , output_type='''np'''						)

							snake_case   :					Tuple    						=				output.images

							snake_case   :					Dict    						=				image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							snake_case   :					Tuple    						=				np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]						)

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



			def      lowerCAmelCase(   self   :  str						):






							"""simple docstring"""
							snake_case   :					Optional[int]    						=				DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''     , subfolder='''scheduler'''						)

							snake_case   :					str    						=				AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''     , scheduler=__UpperCamelCase     , safety_checker=__UpperCamelCase						)
							snake_case   :					List[Any]    						=				alt_pipe.to(__UpperCamelCase						)
							alt_pipe.set_progress_bar_config(disable=__UpperCamelCase						)

							snake_case   :					int    						=				"""A painting of a squirrel eating a burger"""
							snake_case   :					Dict    						=				torch.manual_seed(0						)

							snake_case   :					int    						=				alt_pipe([prompt]     , generator=__UpperCamelCase     , num_inference_steps=2     , output_type='''numpy'''						)
							snake_case   :					List[Any]    						=				output.images

							snake_case   :					Optional[int]    						=				image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							snake_case   :					Optional[int]    						=				np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]						)

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

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

import sentencepiece as spm

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


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


SCREAMING_SNAKE_CASE		:     Union[str, Any]         =	"▁"

SCREAMING_SNAKE_CASE		:     Dict         =	{"vocab_file": "spiece.model"}

SCREAMING_SNAKE_CASE		:     List[Any]         =	{
    "vocab_file": {
        "google/reformer-crime-and-punishment": (
            "https://huggingface.co/google/reformer-crime-and-punishment/resolve/main/spiece.model"
        )
    }
}

SCREAMING_SNAKE_CASE		:     Any         =	{
    "google/reformer-crime-and-punishment": 524288,
}

class        UpperCamelCase						(				__a				):
						a__		:Optional[Any]      =			VOCAB_FILES_NAMES
						a__		:List[Any]      =			PRETRAINED_VOCAB_FILES_MAP
						a__		:Union[str, Any]      =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
						a__		:List[Any]      =			['''input_ids''', '''attention_mask''']





						def __init__(self			,     __UpperCamelCase			,     __UpperCamelCase="</s>"			,     __UpperCamelCase="<unk>"			,     __UpperCamelCase=[]			,     __UpperCamelCase = None			,     **__UpperCamelCase			,     )				->				None:
												UpperCamelCase_		:							List[str]  =		{} if sp_model_kwargs is None else sp_model_kwargs

												super().__init__(
												    eos_token=__UpperCamelCase			,     unk_token=__UpperCamelCase			,     additional_special_tokens=__UpperCamelCase			,     sp_model_kwargs=self.sp_model_kwargs			,     **__UpperCamelCase			,     )

												UpperCamelCase_		:							int  =		vocab_file
												UpperCamelCase_		:							str  =		spm.SentencePieceProcessor(**self.sp_model_kwargs						)
												self.sp_model.Load(__UpperCamelCase						)





						@property
						def        A_							(self						)				->				Dict:
												return self.sp_model.get_piece_size()





						def        A_							(self						)				->				Dict[str, int]:
												UpperCamelCase_		:							str  =		{self.convert_ids_to_tokens(__UpperCamelCase						): i for i in range(self.vocab_size						)}
												vocab.update(self.added_tokens_encoder						)
												return vocab





						def __getstate__(self						)				->				Any:
												UpperCamelCase_		:							Dict  =		self.__dict__.copy()
												UpperCamelCase_		:							List[str]  =		None
												return state





						def __setstate__(self			,     __UpperCamelCase						)				->				List[Any]:
												UpperCamelCase_		:							List[Any]  =		d

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

												UpperCamelCase_		:							Dict  =		spm.SentencePieceProcessor(**self.sp_model_kwargs						)
												self.sp_model.Load(self.vocab_file						)





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





						def        A_							(self			,     __UpperCamelCase						)				->				Union[str, Any]:
												return self.sp_model.piece_to_id(__UpperCamelCase						)





						def        A_							(self			,     __UpperCamelCase						)				->				Dict:
												if index < self.sp_model.get_piece_size():
																		UpperCamelCase_		:							Union[str, Any]  =		self.sp_model.IdToPiece(__UpperCamelCase						)
												return token





						def        A_							(self			,     __UpperCamelCase						)				->				int:
												UpperCamelCase_		:							Dict  =		[]
												UpperCamelCase_		:							Union[str, Any]  =		""""""
												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(__UpperCamelCase						) + token
																								UpperCamelCase_		:							Optional[Any]  =		[]
																		else:
																								current_sub_tokens.append(__UpperCamelCase						)
												out_string += self.sp_model.decode(__UpperCamelCase						)
												return out_string.strip()





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

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

												return (out_vocab_file,)






