


"""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 						snake_case_(				a__							,    a__							,    a__							,    unittest.TestCase	):
     __UpperCamelCase     =    AltDiffusionPipeline
     __UpperCamelCase     =    TEXT_TO_IMAGE_PARAMS
     __UpperCamelCase     =    TEXT_TO_IMAGE_BATCH_PARAMS
     __UpperCamelCase     =    TEXT_TO_IMAGE_IMAGE_PARAMS
     __UpperCamelCase     =    TEXT_TO_IMAGE_IMAGE_PARAMS






     def    lowerCamelCase__   (		self     :    Union[str, Any] ):
           torch.manual_seed(0 )
           lowerCAmelCase    :		List[str]				=      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''')    , cross_attention_dim=3_2    , )
           lowerCAmelCase    :		Optional[Any]				=      DDIMScheduler(
               beta_start=0.00_085    , beta_end=0.012    , beta_schedule='''scaled_linear'''    , clip_sample=UpperCamelCase_    , set_alpha_to_one=UpperCamelCase_    , )
           torch.manual_seed(0 )
           lowerCAmelCase    :		int				=      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    , )

           # 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 )
           lowerCAmelCase    :		List[Any]				=      CLIPTextConfig(
               bos_token_id=0    , eos_token_id=2    , hidden_size=3_2    , projection_dim=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=5_0_0_2    , )
           lowerCAmelCase    :		str				=      CLIPTextModel(UpperCamelCase_ )

           lowerCAmelCase    :		Any				=      XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta''' )
           lowerCAmelCase    :		Union[str, Any]				=      7_7

           lowerCAmelCase    :		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     :    int    , UpperCamelCase_     :    Dict    , UpperCamelCase_     :    List[Any]=0 ):
           if str(UpperCamelCase_ ).startswith('''mps''' ):
                 lowerCAmelCase    :		Optional[int]				=      torch.manual_seed(UpperCamelCase_ )
           else:
                 lowerCAmelCase    :		List[Any]				=      torch.Generator(device=UpperCamelCase_ ).manual_seed(UpperCamelCase_ )
           lowerCAmelCase    :		Optional[Any]				=      {
               '''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     :    Any ):
           super().test_attention_slicing_forward_pass(expected_max_diff=3E-3 )






     def    lowerCamelCase__   (		self     :    Optional[int] ):
           super().test_inference_batch_single_identical(expected_max_diff=3E-3 )






     def    lowerCamelCase__   (		self     :    str ):
           lowerCAmelCase    :		int				=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

           lowerCAmelCase    :		Union[str, Any]				=      self.get_dummy_components()
           torch.manual_seed(0 )
           lowerCAmelCase    :		List[Any]				=      RobertaSeriesConfig(
               hidden_size=3_2    , project_dim=3_2    , intermediate_size=3_7    , layer_norm_eps=1E-05    , num_attention_heads=4    , num_hidden_layers=5    , vocab_size=5_0_0_2    , )
           # TODO: remove after fixing the non-deterministic text encoder
           lowerCAmelCase    :		Optional[Any]				=      RobertaSeriesModelWithTransformation(UpperCamelCase_ )
           lowerCAmelCase    :		Optional[Any]				=      text_encoder

           lowerCAmelCase    :		int				=      AltDiffusionPipeline(**UpperCamelCase_ )
           lowerCAmelCase    :		Any				=      alt_pipe.to(UpperCamelCase_ )
           alt_pipe.set_progress_bar_config(disable=UpperCamelCase_ )

           lowerCAmelCase    :		str				=      self.get_dummy_inputs(UpperCamelCase_ )
           lowerCAmelCase    :		List[str]				=      '''A photo of an astronaut'''
           lowerCAmelCase    :		Optional[Any]				=      alt_pipe(**UpperCamelCase_ )
           lowerCAmelCase    :		Any				=      output.images
           lowerCAmelCase    :		Union[str, Any]				=      image[0, -3:, -3:, -1]

           assert image.shape == (1, 6_4, 6_4, 3)
           lowerCAmelCase    :		Any				=      np.array(
               [0.5_748_162, 0.60_447_145, 0.48_821_217, 0.50_100_636, 0.5_431_185, 0.45_763_683, 0.49_657_696, 0.48_132_733, 0.47_573_093] )

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






     def    lowerCamelCase__   (		self     :    Optional[int] ):
           lowerCAmelCase    :		int				=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

           lowerCAmelCase    :		Dict				=      self.get_dummy_components()
           lowerCAmelCase    :		int				=      PNDMScheduler(skip_prk_steps=UpperCamelCase_ )
           torch.manual_seed(0 )
           lowerCAmelCase    :		List[Any]				=      RobertaSeriesConfig(
               hidden_size=3_2    , project_dim=3_2    , intermediate_size=3_7    , layer_norm_eps=1E-05    , num_attention_heads=4    , num_hidden_layers=5    , vocab_size=5_0_0_2    , )
           # TODO: remove after fixing the non-deterministic text encoder
           lowerCAmelCase    :		str				=      RobertaSeriesModelWithTransformation(UpperCamelCase_ )
           lowerCAmelCase    :		List[str]				=      text_encoder
           lowerCAmelCase    :		Optional[int]				=      AltDiffusionPipeline(**UpperCamelCase_ )
           lowerCAmelCase    :		int				=      alt_pipe.to(UpperCamelCase_ )
           alt_pipe.set_progress_bar_config(disable=UpperCamelCase_ )

           lowerCAmelCase    :		int				=      self.get_dummy_inputs(UpperCamelCase_ )
           lowerCAmelCase    :		Union[str, Any]				=      alt_pipe(**UpperCamelCase_ )
           lowerCAmelCase    :		str				=      output.images
           lowerCAmelCase    :		int				=      image[0, -3:, -3:, -1]

           assert image.shape == (1, 6_4, 6_4, 3)
           lowerCAmelCase    :		Any				=      np.array(
               [0.51_605_093, 0.5_707_241, 0.47_365_507, 0.50_578_886, 0.5_633_877, 0.4_642_503, 0.5_182_081, 0.48_763_484, 0.49_084_237] )

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



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






     def    lowerCamelCase__   (		self     :    int ):
           # clean up the VRAM after each test
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()






     def    lowerCamelCase__   (		self     :    List[Any] ):
           # make sure here that pndm scheduler skips prk
           lowerCAmelCase    :		str				=      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''    , safety_checker=UpperCamelCase_ )
           lowerCAmelCase    :		Any				=      alt_pipe.to(UpperCamelCase_ )
           alt_pipe.set_progress_bar_config(disable=UpperCamelCase_ )

           lowerCAmelCase    :		Tuple				=      '''A painting of a squirrel eating a burger'''
           lowerCAmelCase    :		Dict				=      torch.manual_seed(0 )
           lowerCAmelCase    :		Optional[Any]				=      alt_pipe([prompt]    , generator=UpperCamelCase_    , guidance_scale=6.0    , num_inference_steps=2_0    , output_type='''np''' )

           lowerCAmelCase    :		int				=      output.images

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

           assert image.shape == (1, 5_1_2, 5_1_2, 3)
           lowerCAmelCase    :		List[str]				=      np.array([0.1_010, 0.0_800, 0.0_794, 0.0_885, 0.0_843, 0.0_762, 0.0_769, 0.0_729, 0.0_586] )

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






     def    lowerCamelCase__   (		self     :    List[str] ):
           lowerCAmelCase    :		Union[str, Any]				=      DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''    , subfolder='''scheduler''' )

           lowerCAmelCase    :		Union[str, Any]				=      AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''    , scheduler=UpperCamelCase_    , safety_checker=UpperCamelCase_ )
           lowerCAmelCase    :		str				=      alt_pipe.to(UpperCamelCase_ )
           alt_pipe.set_progress_bar_config(disable=UpperCamelCase_ )

           lowerCAmelCase    :		Dict				=      '''A painting of a squirrel eating a burger'''
           lowerCAmelCase    :		str				=      torch.manual_seed(0 )

           lowerCAmelCase    :		List[str]				=      alt_pipe([prompt]    , generator=UpperCamelCase_    , num_inference_steps=2    , output_type='''numpy''' )
           lowerCAmelCase    :		str				=      output.images

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

           assert image.shape == (1, 5_1_2, 5_1_2, 3)
           lowerCAmelCase    :		Dict				=      np.array([0.4_019, 0.4_052, 0.3_810, 0.4_119, 0.3_916, 0.3_982, 0.4_651, 0.4_195, 0.5_323] )

           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:
 snake_case__						:   str			 = None

snake_case__						:   Optional[Any]			 = logging.get_logger(__name__)
snake_case__						:   Optional[int]			 = {'''vocab_file''': '''spiece.model''', '''tokenizer_file''': '''tokenizer.json'''}

snake_case__						:   Dict			 = {
    '''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''',
    },
}

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

snake_case__						:   Dict			 = '''▁'''



class 						snake_case_(				a__	):
     __UpperCamelCase     =    VOCAB_FILES_NAMES
     __UpperCamelCase     =    PRETRAINED_VOCAB_FILES_MAP
     __UpperCamelCase     =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __UpperCamelCase     =    ['''input_ids''', '''token_type_ids''']
     __UpperCamelCase     =    FNetTokenizer






     def __init__(		self     :    Union[str, Any]    , UpperCamelCase_     :    Union[str, Any]=None    , UpperCamelCase_     :    Union[str, Any]=None    , UpperCamelCase_     :    Any=False    , UpperCamelCase_     :    Any=True    , UpperCamelCase_     :    Dict=True    , UpperCamelCase_     :    Tuple="<unk>"    , UpperCamelCase_     :    List[str]="[SEP]"    , UpperCamelCase_     :    List[Any]="<pad>"    , UpperCamelCase_     :    Union[str, Any]="[CLS]"    , UpperCamelCase_     :    int="[MASK]"    , **UpperCamelCase_     :    Optional[Any]    , ):
           # Mask token behave like a normal word, i.e. include the space before it and
           # is included in the raw text, there should be a match in a non-normalized sentence.
           lowerCAmelCase    :		int				=      (
               AddedToken(UpperCamelCase_    , lstrip=UpperCamelCase_    , rstrip=UpperCamelCase_    , normalized=UpperCamelCase_ )
               if isinstance(UpperCamelCase_    , UpperCamelCase_ )
               else mask_token
           )

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

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






     def    lowerCamelCase__   (		self     :    List[Any]    , UpperCamelCase_     :    List[int]    , UpperCamelCase_     :    Optional[List[int]] = None ):
           lowerCAmelCase    :		Optional[int]				=      [self.sep_token_id]
           lowerCAmelCase    :		Optional[Any]				=      [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     :    List[str]    , UpperCamelCase_     :    List[int]    , UpperCamelCase_     :    Optional[List[int]] = None ):
           lowerCAmelCase    :		List[str]				=      [self.sep_token_id]
           lowerCAmelCase    :		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     :    List[Any]    , UpperCamelCase_     :    str    , UpperCamelCase_     :    Optional[str] = None ):
           if not os.path.isdir(UpperCamelCase_ ):
                 logger.error(F'''Vocabulary path ({save_directory}) should be a directory''' )
                 return
           lowerCAmelCase    :		str				=      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_ ):
                 copyfile(self.vocab_file    , UpperCamelCase_ )

           return (out_vocab_file,)


