






import os
import tempfile
import unittest

import numpy as np

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax, slow


if is_flax_available():
							import jax
							import jax.numpy as jnp
							from flax.jax_utils import replicate
							from flax.training.common_utils import shard

							from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline

@require_flax
class 				lowerCAmelCase_	(		unittest.TestCase			):


				def 						__a					(					self      ):
											with tempfile.TemporaryDirectory() as tmpdirname:
																		# pipeline has Flax weights
																		_lowercase							:   Optional[Any] 							=   FlaxDiffusionPipeline.from_pretrained(
																		    'hf-internal-testing/tiny-stable-diffusion-pipe' ,						safety_checker=_lowerCAmelCase ,						cache_dir=_lowerCAmelCase      )

																		_lowercase							:   str 							=   [t[-1] for t in os.walk(os.path.join(_lowerCAmelCase ,						os.listdir(_lowerCAmelCase      )[0] ,						'snapshots'      )      )]
																		_lowercase							:   Tuple 							=   [item for sublist in all_root_files for item in sublist]

																		# None of the downloaded files should be a PyTorch file even if we have some here:
																		# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
																		assert not any(f.endswith('.bin'      ) for f in files      )




@slow
@require_flax
class 				lowerCAmelCase_	(		unittest.TestCase			):


				def 						__a					(					self      ):
											_lowercase							,							_lowercase							:   Optional[Any] 							=   FlaxStableDiffusionPipeline.from_pretrained(
											    'hf-internal-testing/tiny-stable-diffusion-pipe' ,						safety_checker=_lowerCAmelCase      )

											_lowercase							:   int 							=   (
											    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
											    ' field, close up, split lighting, cinematic'
											)

											_lowercase							:   Optional[Any] 							=   jax.random.PRNGKey(0      )
											_lowercase							:   Optional[Any] 							=   4

											_lowercase							:   Any 							=   jax.device_count()
											_lowercase							:   Dict 							=   num_samples * [prompt]
											_lowercase							:   List[str] 							=   pipeline.prepare_inputs(_lowerCAmelCase      )

											# shard inputs and rng
											_lowercase							:   Tuple 							=   replicate(_lowerCAmelCase      )
											_lowercase							:   Optional[Any] 							=   jax.random.split(_lowerCAmelCase ,						_lowerCAmelCase      )
											_lowercase							:   List[str] 							=   shard(_lowerCAmelCase      )

											_lowercase							:   List[Any] 							=   pipeline(_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						jit=_lowerCAmelCase      ).images

											assert images.shape == (num_samples, 1, 6_4, 6_4, 3)
											if jax.device_count() == 8:
																		assert np.abs(np.abs(images[0, 0, :2, :2, -2:] ,						dtype=np.floataa      ).sum() - 4.1_51_47_45      ) < 1E-3
																		assert np.abs(np.abs(_lowerCAmelCase ,						dtype=np.floataa      ).sum() - 4_99_47.8_75      ) < 5E-1

											_lowercase							:   Any 							=   pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]      )      )      )
											assert len(_lowerCAmelCase      ) == num_samples


				def 						__a					(					self      ):
											_lowercase							,							_lowercase							:   Optional[Any] 							=   FlaxStableDiffusionPipeline.from_pretrained(
											    'CompVis/stable-diffusion-v1-4' ,						revision='flax' ,						safety_checker=_lowerCAmelCase      )

											_lowercase							:   Dict 							=   (
											    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
											    ' field, close up, split lighting, cinematic'
											)

											_lowercase							:   Tuple 							=   jax.random.PRNGKey(0      )
											_lowercase							:   Optional[int] 							=   5_0

											_lowercase							:   Optional[Any] 							=   jax.device_count()
											_lowercase							:   List[str] 							=   num_samples * [prompt]
											_lowercase							:   Dict 							=   pipeline.prepare_inputs(_lowerCAmelCase      )

											# shard inputs and rng
											_lowercase							:   str 							=   replicate(_lowerCAmelCase      )
											_lowercase							:   Optional[int] 							=   jax.random.split(_lowerCAmelCase ,						_lowerCAmelCase      )
											_lowercase							:   Optional[Any] 							=   shard(_lowerCAmelCase      )

											_lowercase							:   Dict 							=   pipeline(_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						jit=_lowerCAmelCase      ).images

											assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
											if jax.device_count() == 8:
																		assert np.abs((np.abs(images[0, 0, :2, :2, -2:] ,						dtype=np.floataa      ).sum() - 0.05_65_24_01)      ) < 1E-3
																		assert np.abs((np.abs(_lowerCAmelCase ,						dtype=np.floataa      ).sum() - 2_38_38_08.2)      ) < 5E-1


				def 						__a					(					self      ):
											_lowercase							,							_lowercase							:   Tuple 							=   FlaxStableDiffusionPipeline.from_pretrained(
											    'CompVis/stable-diffusion-v1-4' ,						revision='bf16' ,						dtype=jnp.bfloataa ,						safety_checker=_lowerCAmelCase      )

											_lowercase							:   Union[str, Any] 							=   (
											    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
											    ' field, close up, split lighting, cinematic'
											)

											_lowercase							:   Any 							=   jax.random.PRNGKey(0      )
											_lowercase							:   Optional[int] 							=   5_0

											_lowercase							:   Optional[Any] 							=   jax.device_count()
											_lowercase							:   Union[str, Any] 							=   num_samples * [prompt]
											_lowercase							:   Any 							=   pipeline.prepare_inputs(_lowerCAmelCase      )

											# shard inputs and rng
											_lowercase							:   Optional[Any] 							=   replicate(_lowerCAmelCase      )
											_lowercase							:   Any 							=   jax.random.split(_lowerCAmelCase ,						_lowerCAmelCase      )
											_lowercase							:   Dict 							=   shard(_lowerCAmelCase      )

											_lowercase							:   Union[str, Any] 							=   pipeline(_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						jit=_lowerCAmelCase      ).images

											assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
											if jax.device_count() == 8:
																		assert np.abs((np.abs(images[0, 0, :2, :2, -2:] ,						dtype=np.floataa      ).sum() - 0.04_00_39_06)      ) < 1E-3
																		assert np.abs((np.abs(_lowerCAmelCase ,						dtype=np.floataa      ).sum() - 2_37_35_16.75)      ) < 5E-1


				def 						__a					(					self      ):
											_lowercase							,							_lowercase							:   Optional[Any] 							=   FlaxStableDiffusionPipeline.from_pretrained(
											    'CompVis/stable-diffusion-v1-4' ,						revision='bf16' ,						dtype=jnp.bfloataa      )

											_lowercase							:   List[Any] 							=   (
											    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
											    ' field, close up, split lighting, cinematic'
											)

											_lowercase							:   Union[str, Any] 							=   jax.random.PRNGKey(0      )
											_lowercase							:   Optional[int] 							=   5_0

											_lowercase							:   List[str] 							=   jax.device_count()
											_lowercase							:   List[Any] 							=   num_samples * [prompt]
											_lowercase							:   List[str] 							=   pipeline.prepare_inputs(_lowerCAmelCase      )

											# shard inputs and rng
											_lowercase							:   Union[str, Any] 							=   replicate(_lowerCAmelCase      )
											_lowercase							:   List[str] 							=   jax.random.split(_lowerCAmelCase ,						_lowerCAmelCase      )
											_lowercase							:   Union[str, Any] 							=   shard(_lowerCAmelCase      )

											_lowercase							:   Dict 							=   pipeline(_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						jit=_lowerCAmelCase      ).images

											assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
											if jax.device_count() == 8:
																		assert np.abs((np.abs(images[0, 0, :2, :2, -2:] ,						dtype=np.floataa      ).sum() - 0.04_00_39_06)      ) < 1E-3
																		assert np.abs((np.abs(_lowerCAmelCase ,						dtype=np.floataa      ).sum() - 2_37_35_16.75)      ) < 5E-1


				def 						__a					(					self      ):
											_lowercase							:   str 							=   FlaxDDIMScheduler(
											    beta_start=0.0_00_85 ,						beta_end=0.0_12 ,						beta_schedule='scaled_linear' ,						set_alpha_to_one=_lowerCAmelCase ,						steps_offset=1 ,						)

											_lowercase							,							_lowercase							:   Optional[int] 							=   FlaxStableDiffusionPipeline.from_pretrained(
											    'CompVis/stable-diffusion-v1-4' ,						revision='bf16' ,						dtype=jnp.bfloataa ,						scheduler=_lowerCAmelCase ,						safety_checker=_lowerCAmelCase ,						)
											_lowercase							:   Any 							=   scheduler.create_state()

											_lowercase							:   List[str] 							=   scheduler_state

											_lowercase							:   Dict 							=   (
											    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
											    ' field, close up, split lighting, cinematic'
											)

											_lowercase							:   List[str] 							=   jax.random.PRNGKey(0      )
											_lowercase							:   int 							=   5_0

											_lowercase							:   List[str] 							=   jax.device_count()
											_lowercase							:   Any 							=   num_samples * [prompt]
											_lowercase							:   int 							=   pipeline.prepare_inputs(_lowerCAmelCase      )

											# shard inputs and rng
											_lowercase							:   str 							=   replicate(_lowerCAmelCase      )
											_lowercase							:   str 							=   jax.random.split(_lowerCAmelCase ,						_lowerCAmelCase      )
											_lowercase							:   str 							=   shard(_lowerCAmelCase      )

											_lowercase							:   Dict 							=   pipeline(_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						jit=_lowerCAmelCase      ).images

											assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
											if jax.device_count() == 8:
																		assert np.abs((np.abs(images[0, 0, :2, :2, -2:] ,						dtype=np.floataa      ).sum() - 0.0_45_04_39_45)      ) < 1E-3
																		assert np.abs((np.abs(_lowerCAmelCase ,						dtype=np.floataa      ).sum() - 2_34_76_93.5)      ) < 5E-1


				def 						__a					(					self      ):
											_lowercase							:   List[Any] 							=   (
											    'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
											    ' field, close up, split lighting, cinematic'
											)

											_lowercase							:   Optional[Any] 							=   jax.device_count()
											_lowercase							:   Union[str, Any] 							=   num_samples * [prompt]
											_lowercase							:   Optional[Any] 							=   jax.random.split(jax.random.PRNGKey(0      ) ,						_lowerCAmelCase      )

											_lowercase							,							_lowercase							:   Optional[Any] 							=   FlaxStableDiffusionPipeline.from_pretrained(
											    'CompVis/stable-diffusion-v1-4' ,						revision='bf16' ,						dtype=jnp.bfloataa ,						safety_checker=_lowerCAmelCase ,						)

											_lowercase							:   Optional[Any] 							=   replicate(_lowerCAmelCase      )
											_lowercase							:   Union[str, Any] 							=   pipeline.prepare_inputs(_lowerCAmelCase      )
											_lowercase							:   Optional[Any] 							=   shard(_lowerCAmelCase      )
											_lowercase							:   Tuple 							=   pipeline(_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						jit=_lowerCAmelCase      ).images
											assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
											_lowercase							:   Any 							=   images[2, 0, 2_5_6, 1_0:1_7, 1]

											# With memory efficient attention
											_lowercase							,							_lowercase							:   Optional[Any] 							=   FlaxStableDiffusionPipeline.from_pretrained(
											    'CompVis/stable-diffusion-v1-4' ,						revision='bf16' ,						dtype=jnp.bfloataa ,						safety_checker=_lowerCAmelCase ,						use_memory_efficient_attention=_lowerCAmelCase ,						)

											_lowercase							:   Optional[Any] 							=   replicate(_lowerCAmelCase      )
											_lowercase							:   str 							=   pipeline.prepare_inputs(_lowerCAmelCase      )
											_lowercase							:   int 							=   shard(_lowerCAmelCase      )
											_lowercase							:   List[Any] 							=   pipeline(_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						jit=_lowerCAmelCase      ).images
											assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
											_lowercase							:   Dict 							=   images[2, 0, 2_5_6, 1_0:1_7, 1]

											# I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
											# over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
											assert abs(slice_eff - slice      ).max() < 1E-2






"""simple docstring"""






import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_funnel import FunnelTokenizer


__SCREAMING_SNAKE_CASE							=     logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE							=     {'vocab_file': 'vocab.txt', 'tokenizer_file': 'tokenizer.json'}

__SCREAMING_SNAKE_CASE							=     [
    'small',
    'small-base',
    'medium',
    'medium-base',
    'intermediate',
    'intermediate-base',
    'large',
    'large-base',
    'xlarge',
    'xlarge-base',
]

__SCREAMING_SNAKE_CASE							=     {
    'vocab_file': {
        'funnel-transformer/small': 'https://huggingface.co/funnel-transformer/small/resolve/main/vocab.txt',
        'funnel-transformer/small-base': 'https://huggingface.co/funnel-transformer/small-base/resolve/main/vocab.txt',
        'funnel-transformer/medium': 'https://huggingface.co/funnel-transformer/medium/resolve/main/vocab.txt',
        'funnel-transformer/medium-base': (
            'https://huggingface.co/funnel-transformer/medium-base/resolve/main/vocab.txt'
        ),
        'funnel-transformer/intermediate': (
            'https://huggingface.co/funnel-transformer/intermediate/resolve/main/vocab.txt'
        ),
        'funnel-transformer/intermediate-base': (
            'https://huggingface.co/funnel-transformer/intermediate-base/resolve/main/vocab.txt'
        ),
        'funnel-transformer/large': 'https://huggingface.co/funnel-transformer/large/resolve/main/vocab.txt',
        'funnel-transformer/large-base': 'https://huggingface.co/funnel-transformer/large-base/resolve/main/vocab.txt',
        'funnel-transformer/xlarge': 'https://huggingface.co/funnel-transformer/xlarge/resolve/main/vocab.txt',
        'funnel-transformer/xlarge-base': (
            'https://huggingface.co/funnel-transformer/xlarge-base/resolve/main/vocab.txt'
        ),
    },
    'tokenizer_file': {
        'funnel-transformer/small': 'https://huggingface.co/funnel-transformer/small/resolve/main/tokenizer.json',
        'funnel-transformer/small-base': (
            'https://huggingface.co/funnel-transformer/small-base/resolve/main/tokenizer.json'
        ),
        'funnel-transformer/medium': 'https://huggingface.co/funnel-transformer/medium/resolve/main/tokenizer.json',
        'funnel-transformer/medium-base': (
            'https://huggingface.co/funnel-transformer/medium-base/resolve/main/tokenizer.json'
        ),
        'funnel-transformer/intermediate': (
            'https://huggingface.co/funnel-transformer/intermediate/resolve/main/tokenizer.json'
        ),
        'funnel-transformer/intermediate-base': (
            'https://huggingface.co/funnel-transformer/intermediate-base/resolve/main/tokenizer.json'
        ),
        'funnel-transformer/large': 'https://huggingface.co/funnel-transformer/large/resolve/main/tokenizer.json',
        'funnel-transformer/large-base': (
            'https://huggingface.co/funnel-transformer/large-base/resolve/main/tokenizer.json'
        ),
        'funnel-transformer/xlarge': 'https://huggingface.co/funnel-transformer/xlarge/resolve/main/tokenizer.json',
        'funnel-transformer/xlarge-base': (
            'https://huggingface.co/funnel-transformer/xlarge-base/resolve/main/tokenizer.json'
        ),
    },
}
__SCREAMING_SNAKE_CASE							=     {F"""funnel-transformer/{name}""": 512 for name in _model_names}
__SCREAMING_SNAKE_CASE							=     {F"""funnel-transformer/{name}""": {'do_lower_case': True} for name in _model_names}





class 				a__					(						A__       ):
     UpperCAmelCase__         =      VOCAB_FILES_NAMES
     UpperCAmelCase__         =      PRETRAINED_VOCAB_FILES_MAP
     UpperCAmelCase__         =      PRETRAINED_INIT_CONFIGURATION
     UpperCAmelCase__         =      FunnelTokenizer
     UpperCAmelCase__         =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     UpperCAmelCase__         =      2



     def __init__(							self :str      ,  _lowerCamelCase :int=None      ,  _lowerCamelCase :List[Any]=None      ,  _lowerCamelCase :Dict=True      ,  _lowerCamelCase :Optional[Any]="<unk>"      ,  _lowerCamelCase :Optional[Any]="<sep>"      ,  _lowerCamelCase :List[Any]="<pad>"      ,  _lowerCamelCase :List[Any]="<cls>"      ,  _lowerCamelCase :Union[str, Any]="<mask>"      ,  _lowerCamelCase :Any="<s>"      ,  _lowerCamelCase :Union[str, Any]="</s>"      ,  _lowerCamelCase :str=True      ,  _lowerCamelCase :Dict=True      ,  _lowerCamelCase :Tuple=None      ,  _lowerCamelCase :Tuple="##"      ,  **_lowerCamelCase :int      ,  ):







      '''simple docstring'''



      super().__init__(
          _lowerCamelCase      ,  tokenizer_file=_lowerCamelCase      ,  do_lower_case=_lowerCamelCase      ,  unk_token=_lowerCamelCase      ,  sep_token=_lowerCamelCase      ,  pad_token=_lowerCamelCase      ,  cls_token=_lowerCamelCase      ,  mask_token=_lowerCamelCase      ,  bos_token=_lowerCamelCase      ,  eos_token=_lowerCamelCase      ,  clean_text=_lowerCamelCase      ,  tokenize_chinese_chars=_lowerCamelCase      ,  strip_accents=_lowerCamelCase      ,  wordpieces_prefix=_lowerCamelCase      ,  **_lowerCamelCase      ,  )

      UpperCamelCase_     :					str   	=json.loads(self.backend_tokenizer.normalizer.__getstate__()   )
      if (
          normalizer_state.get('lowercase'      ,  _lowerCamelCase   ) != do_lower_case
          or normalizer_state.get('strip_accents'      ,  _lowerCamelCase   ) != strip_accents
          or normalizer_state.get('handle_chinese_chars'      ,  _lowerCamelCase   ) != tokenize_chinese_chars
      ):
       UpperCamelCase_     :					int   	=getattr(_lowerCamelCase      ,  normalizer_state.pop('type'   )   )
       UpperCamelCase_     :					Dict   	=do_lower_case
       UpperCamelCase_     :					Tuple   	=strip_accents
       UpperCamelCase_     :					int   	=tokenize_chinese_chars
       UpperCamelCase_     :					Tuple   	=normalizer_class(**_lowerCamelCase   )

      UpperCamelCase_     :					Dict   	=do_lower_case



     def  lowerCamelCase_						(							self :Optional[int]      ,  _lowerCamelCase :str      ,  _lowerCamelCase :Optional[int]=None   ):







      '''simple docstring'''



      UpperCamelCase_     :					Optional[Any]   	=[self.cls_token_id] + token_ids_a + [self.sep_token_id]

      if token_ids_a:
       output += token_ids_a + [self.sep_token_id]

      return output



     def  lowerCamelCase_						(							self :List[Any]      ,  _lowerCamelCase :List[int]      ,  _lowerCamelCase :Optional[List[int]] = None   ):







      '''simple docstring'''



      UpperCamelCase_     :					Optional[int]   	=[self.sep_token_id]
      UpperCamelCase_     :					List[Any]   	=[self.cls_token_id]
      if token_ids_a is None:
       return len(cls   ) * [self.cls_token_type_id] + len(token_ids_a + sep   ) * [0]
      return len(cls   ) * [self.cls_token_type_id] + len(token_ids_a + sep   ) * [0] + len(token_ids_a + sep   ) * [1]




     def  lowerCamelCase_						(							self :Optional[Any]      ,  _lowerCamelCase :str      ,  _lowerCamelCase :Optional[str] = None   ):







      '''simple docstring'''



      UpperCamelCase_     :					int   	=self._tokenizer.model.save(_lowerCamelCase      ,  name=_lowerCamelCase   )
      return tuple(_lowerCamelCase   )
