


'''simple docstring'''


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 		lowerCAmelCase__	(	UpperCAmelCase__			,				UpperCAmelCase__			,				UpperCAmelCase__			,				unittest.TestCase		):
			lowerCAmelCase       :		Optional[Any]				      =   AltDiffusionPipeline
			lowerCAmelCase       :		Tuple				      =   TEXT_TO_IMAGE_PARAMS
			lowerCAmelCase       :		Optional[Any]				      =   TEXT_TO_IMAGE_BATCH_PARAMS
			lowerCAmelCase       :		Dict				      =   TEXT_TO_IMAGE_IMAGE_PARAMS
			lowerCAmelCase       :		Union[str, Any]				      =   TEXT_TO_IMAGE_IMAGE_PARAMS




			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->Any:



					'''simple docstring'''







					torch.manual_seed(0       )
					_UpperCAmelCase						:				List[str]	    =					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		,     )
					_UpperCAmelCase						:				Dict	    =					DDIMScheduler(
					    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						:				Optional[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       )
					_UpperCAmelCase						:				Tuple	    =					CLIPTextConfig(
					    bos_token_id=0		,     eos_token_id=2		,     hidden_size=32		,     projection_dim=32		,     intermediate_size=37		,     layer_norm_eps=1E-05		,     num_attention_heads=4		,     num_hidden_layers=5		,     pad_token_id=1		,     vocab_size=50_02		,     )
					_UpperCAmelCase						:				Optional[Any]	    =					CLIPTextModel(lowerCamelCase__       )

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

					_UpperCAmelCase						:				Optional[Any]	    =					{
					    "unet": unet,
					    "scheduler": scheduler,
					    "vae": vae,
					    "text_encoder": text_encoder,
					    "tokenizer": tokenizer,
					    "safety_checker": None,
					    "feature_extractor": None,
					}
					return components




			def 							lowerCAmelCase__			(						self     :   Union[str, Any]		,     lowerCamelCase__     :   Optional[Any]		,     lowerCamelCase__     :   List[Any]=0       )      ->Union[str, Any]:



					'''simple docstring'''







					if str(lowerCamelCase__       ).startswith("mps"       ):
							_UpperCAmelCase						:				str	    =					torch.manual_seed(lowerCamelCase__       )
					else:
							_UpperCAmelCase						:				str	    =					torch.Generator(device=lowerCamelCase__       ).manual_seed(lowerCamelCase__       )
					_UpperCAmelCase						:				int	    =					{
					    "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     :   List[str]       )      ->Union[str, Any]:



					'''simple docstring'''







					super().test_attention_slicing_forward_pass(expected_max_diff=3E-3       )




			def 							lowerCAmelCase__			(						self     :   Dict       )      ->List[str]:



					'''simple docstring'''







					super().test_inference_batch_single_identical(expected_max_diff=3E-3       )




			def 							lowerCAmelCase__			(						self     :   int       )      ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase						:				Any	    =					"cpu"  # ensure determinism for the device-dependent torch.Generator

					_UpperCAmelCase						:				Any	    =					self.get_dummy_components()
					torch.manual_seed(0       )
					_UpperCAmelCase						:				Dict	    =					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=50_02		,     )
					# TODO: remove after fixing the non-deterministic text encoder
					_UpperCAmelCase						:				List[Any]	    =					RobertaSeriesModelWithTransformation(lowerCamelCase__       )
					_UpperCAmelCase						:				Optional[int]	    =					text_encoder

					_UpperCAmelCase						:				str	    =					AltDiffusionPipeline(**lowerCamelCase__       )
					_UpperCAmelCase						:				List[Any]	    =					alt_pipe.to(lowerCamelCase__       )
					alt_pipe.set_progress_bar_config(disable=lowerCamelCase__       )

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

					assert image.shape == (1, 64, 64, 3)
					_UpperCAmelCase						:				List[str]	    =					np.array(
					    [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]       )

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






			def 							lowerCAmelCase__			(						self     :   Dict       )      ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[int]	    =					"cpu"  # ensure determinism for the device-dependent torch.Generator

					_UpperCAmelCase						:				Any	    =					self.get_dummy_components()
					_UpperCAmelCase						:				List[str]	    =					PNDMScheduler(skip_prk_steps=lowerCamelCase__       )
					torch.manual_seed(0       )
					_UpperCAmelCase						:				List[str]	    =					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=50_02		,     )
					# TODO: remove after fixing the non-deterministic text encoder
					_UpperCAmelCase						:				str	    =					RobertaSeriesModelWithTransformation(lowerCamelCase__       )
					_UpperCAmelCase						:				Optional[int]	    =					text_encoder
					_UpperCAmelCase						:				Optional[int]	    =					AltDiffusionPipeline(**lowerCamelCase__       )
					_UpperCAmelCase						:				Optional[Any]	    =					alt_pipe.to(lowerCamelCase__       )
					alt_pipe.set_progress_bar_config(disable=lowerCamelCase__       )

					_UpperCAmelCase						:				Tuple	    =					self.get_dummy_inputs(lowerCamelCase__       )
					_UpperCAmelCase						:				Dict	    =					alt_pipe(**lowerCamelCase__       )
					_UpperCAmelCase						:				Optional[int]	    =					output.images
					_UpperCAmelCase						:				Optional[int]	    =					image[0, -3:, -3:, -1]

					assert image.shape == (1, 64, 64, 3)
					_UpperCAmelCase						:				Union[str, Any]	    =					np.array(
					    [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]       )

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






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




			def 							lowerCAmelCase__			(						self     :   Tuple       )      ->Union[str, Any]:



					'''simple docstring'''







					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()




			def 							lowerCAmelCase__			(						self     :   Optional[Any]       )      ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase						:				Any	    =					AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"		,     safety_checker=lowerCamelCase__       )
					_UpperCAmelCase						:				int	    =					alt_pipe.to(lowerCamelCase__       )
					alt_pipe.set_progress_bar_config(disable=lowerCamelCase__       )

					_UpperCAmelCase						:				int	    =					"A painting of a squirrel eating a burger"
					_UpperCAmelCase						:				Any	    =					torch.manual_seed(0       )
					_UpperCAmelCase						:				Any	    =					alt_pipe([prompt]		,     generator=lowerCamelCase__		,     guidance_scale=6.0		,     num_inference_steps=20		,     output_type="np"       )

					_UpperCAmelCase						:				List[str]	    =					output.images

					_UpperCAmelCase						:				Any	    =					image[0, -3:, -3:, -1]

					assert image.shape == (1, 5_12, 5_12, 3)
					_UpperCAmelCase						:				Dict	    =					np.array([0.1_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]       )

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






			def 							lowerCAmelCase__			(						self     :   Dict       )      ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				str	    =					DDIMScheduler.from_pretrained("BAAI/AltDiffusion"		,     subfolder="scheduler"       )

					_UpperCAmelCase						:				str	    =					AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"		,     scheduler=lowerCamelCase__		,     safety_checker=lowerCamelCase__       )
					_UpperCAmelCase						:				List[str]	    =					alt_pipe.to(lowerCamelCase__       )
					alt_pipe.set_progress_bar_config(disable=lowerCamelCase__       )

					_UpperCAmelCase						:				int	    =					"A painting of a squirrel eating a burger"
					_UpperCAmelCase						:				Optional[Any]	    =					torch.manual_seed(0       )

					_UpperCAmelCase						:				int	    =					alt_pipe([prompt]		,     generator=lowerCamelCase__		,     num_inference_steps=2		,     output_type="numpy"       )
					_UpperCAmelCase						:				str	    =					output.images

					_UpperCAmelCase						:				List[Any]	    =					image[0, -3:, -3:, -1]

					assert image.shape == (1, 5_12, 5_12, 3)
					_UpperCAmelCase						:				Dict	    =					np.array([0.4_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]       )

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







'''simple docstring'''


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

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


if is_sentencepiece_available():
						from .tokenization_fnet import FNetTokenizer
else:
						lowerCamelCase__					=   None

lowerCamelCase__					=   logging.get_logger(__name__)
lowerCamelCase__					=   {'vocab_file': 'spiece.model', 'tokenizer_file': 'tokenizer.json'}

lowerCamelCase__					=   {
    'vocab_file': {
        'google/fnet-base': 'https://huggingface.co/google/fnet-base/resolve/main/spiece.model',
        'google/fnet-large': 'https://huggingface.co/google/fnet-large/resolve/main/spiece.model',
    },
    'tokenizer_file': {
        'google/fnet-base': 'https://huggingface.co/google/fnet-base/resolve/main/tokenizer.json',
        'google/fnet-large': 'https://huggingface.co/google/fnet-large/resolve/main/tokenizer.json',
    },
}

lowerCamelCase__					=   {
    'google/fnet-base': 512,
    'google/fnet-large': 512,
}

lowerCamelCase__					=   '▁'






class 		lowerCAmelCase__	(	UpperCAmelCase__		):
			lowerCAmelCase       :		Dict				      =   VOCAB_FILES_NAMES
			lowerCAmelCase       :		Optional[Any]				      =   PRETRAINED_VOCAB_FILES_MAP
			lowerCAmelCase       :		List[str]				      =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			lowerCAmelCase       :		Optional[int]				      =   ["input_ids", "token_type_ids"]
			lowerCAmelCase       :		Optional[Any]				      =   FNetTokenizer




			def __init__(						self     :   Dict		,     lowerCamelCase__     :   int=None		,     lowerCamelCase__     :   List[Any]=None		,     lowerCamelCase__     :   Any=False		,     lowerCamelCase__     :   List[str]=True		,     lowerCamelCase__     :   Dict=True		,     lowerCamelCase__     :   str="<unk>"		,     lowerCamelCase__     :   List[str]="[SEP]"		,     lowerCamelCase__     :   Union[str, Any]="<pad>"		,     lowerCamelCase__     :   Optional[Any]="[CLS]"		,     lowerCamelCase__     :   Any="[MASK]"		,     **lowerCamelCase__     :   Any		,     )      ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				Dict	    =					(
					    AddedToken(lowerCamelCase__		,     lstrip=lowerCamelCase__		,     rstrip=lowerCamelCase__		,     normalized=lowerCamelCase__       )
					    if isinstance(lowerCamelCase__		,     lowerCamelCase__       )
					    else mask_token
					)

					super().__init__(
					    lowerCamelCase__		,     tokenizer_file=lowerCamelCase__		,     do_lower_case=lowerCamelCase__		,     remove_space=lowerCamelCase__		,     keep_accents=lowerCamelCase__		,     unk_token=lowerCamelCase__		,     sep_token=lowerCamelCase__		,     pad_token=lowerCamelCase__		,     cls_token=lowerCamelCase__		,     mask_token=lowerCamelCase__		,     **lowerCamelCase__		,     )

					_UpperCAmelCase						:				Optional[Any]	    =					do_lower_case
					_UpperCAmelCase						:				Tuple	    =					remove_space
					_UpperCAmelCase						:				List[Any]	    =					keep_accents
					_UpperCAmelCase						:				Tuple	    =					vocab_file
					_UpperCAmelCase						:				str	    =					False if not self.vocab_file else True




			def 							lowerCAmelCase__			(						self     :   str		,     lowerCamelCase__     :   List[int]		,     lowerCamelCase__     :   Optional[List[int]] = None       )      ->List[int]:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[Any]	    =					[self.sep_token_id]
					_UpperCAmelCase						:				Tuple	    =					[self.cls_token_id]
					if token_ids_a is None:
							return cls + token_ids_a + sep
					return cls + token_ids_a + sep + token_ids_a + sep




			def 							lowerCAmelCase__			(						self     :   int		,     lowerCamelCase__     :   List[int]		,     lowerCamelCase__     :   Optional[List[int]] = None       )      ->List[int]:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[int]	    =					[self.sep_token_id]
					_UpperCAmelCase						:				Optional[Any]	    =					[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]






			def 							lowerCAmelCase__			(						self     :   Dict		,     lowerCamelCase__     :   str		,     lowerCamelCase__     :   Optional[str] = None       )      ->Tuple[str]:



					'''simple docstring'''







					if not os.path.isdir(lowerCamelCase__       ):
							logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""       )
							return
					_UpperCAmelCase						:				Union[str, 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,)



