






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  ):
								"""simple docstring"""
								with tempfile.TemporaryDirectory() as tmpdirname:
												# pipeline has Flax weights
												__magic_name__		:Optional[Any]   =   FlaxDiffusionPipeline.from_pretrained(
												    '''hf-internal-testing/tiny-stable-diffusion-pipe'''							,       safety_checker=__lowerCAmelCase							,       cache_dir=__lowerCAmelCase  )

												__magic_name__		:Dict   =   [t[-1] for t in os.walk(os.path.join(__lowerCAmelCase							,       os.listdir(__lowerCAmelCase  )[0]							,       '''snapshots'''  )  )]
												__magic_name__		: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  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:Tuple   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''hf-internal-testing/tiny-stable-diffusion-pipe'''							,       safety_checker=__lowerCAmelCase  )

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

								__magic_name__		:int   =   jax.random.PRNGKey(0  )
								__magic_name__		:Any   =   4

								__magic_name__		:str   =   jax.device_count()
								__magic_name__		:List[str]   =   num_samples * [prompt]
								__magic_name__		:str   =   pipeline.prepare_inputs(__lowerCAmelCase  )

								# shard inputs and rng
								__magic_name__		:List[Any]   =   replicate(__lowerCAmelCase  )
								__magic_name__		:int   =   jax.random.split(__lowerCAmelCase							,       __lowerCAmelCase  )
								__magic_name__		:int   =   shard(__lowerCAmelCase  )

								__magic_name__		: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.1514745  ) < 1E-3
												assert np.abs(np.abs(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 49947.875  ) < 5E-1

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







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:str   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''							,       revision='''flax'''							,       safety_checker=__lowerCAmelCase  )

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

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

								__magic_name__		:Tuple   =   jax.device_count()
								__magic_name__		:Any   =   num_samples * [prompt]
								__magic_name__		:Dict   =   pipeline.prepare_inputs(__lowerCAmelCase  )

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

								__magic_name__		:Optional[int]   =   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.05652401)  ) < 1E-3
												assert np.abs((np.abs(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 2383808.2)  ) < 5E-1







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:Union[str, Any]   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''							,       revision='''bf16'''							,       dtype=jnp.bfloataa							,       safety_checker=__lowerCAmelCase  )

								__magic_name__		: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'''
								)

								__magic_name__		:Tuple   =   jax.random.PRNGKey(0  )
								__magic_name__		:str   =   5_0

								__magic_name__		:Optional[Any]   =   jax.device_count()
								__magic_name__		:List[Any]   =   num_samples * [prompt]
								__magic_name__		:List[Any]   =   pipeline.prepare_inputs(__lowerCAmelCase  )

								# shard inputs and rng
								__magic_name__		:Dict   =   replicate(__lowerCAmelCase  )
								__magic_name__		:Tuple   =   jax.random.split(__lowerCAmelCase							,       __lowerCAmelCase  )
								__magic_name__		:List[Any]   =   shard(__lowerCAmelCase  )

								__magic_name__		:List[str]   =   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.04003906)  ) < 1E-3
												assert np.abs((np.abs(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 2373516.75)  ) < 5E-1







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__  ,					__magic_name__		:Dict   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''							,       revision='''bf16'''							,       dtype=jnp.bfloataa  )

								__magic_name__		: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'''
								)

								__magic_name__		:Any   =   jax.random.PRNGKey(0  )
								__magic_name__		:Any   =   5_0

								__magic_name__		:Optional[Any]   =   jax.device_count()
								__magic_name__		:Optional[Any]   =   num_samples * [prompt]
								__magic_name__		:Optional[Any]   =   pipeline.prepare_inputs(__lowerCAmelCase  )

								# shard inputs and rng
								__magic_name__		:Tuple   =   replicate(__lowerCAmelCase  )
								__magic_name__		:Union[str, Any]   =   jax.random.split(__lowerCAmelCase							,       __lowerCAmelCase  )
								__magic_name__		:Tuple   =   shard(__lowerCAmelCase  )

								__magic_name__		:Tuple   =   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.04003906)  ) < 1E-3
												assert np.abs((np.abs(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 2373516.75)  ) < 5E-1







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   FlaxDDIMScheduler(
								    beta_start=0.00085							,       beta_end=0.012							,       beta_schedule='''scaled_linear'''							,       set_alpha_to_one=__lowerCAmelCase							,       steps_offset=1							,       )

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

								__magic_name__		:Optional[int]   =   scheduler_state

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

								__magic_name__		:Tuple   =   jax.random.PRNGKey(0  )
								__magic_name__		:Any   =   5_0

								__magic_name__		:Dict   =   jax.device_count()
								__magic_name__		:List[Any]   =   num_samples * [prompt]
								__magic_name__		:Optional[int]   =   pipeline.prepare_inputs(__lowerCAmelCase  )

								# shard inputs and rng
								__magic_name__		:List[str]   =   replicate(__lowerCAmelCase  )
								__magic_name__		:Tuple   =   jax.random.split(__lowerCAmelCase							,       __lowerCAmelCase  )
								__magic_name__		:int   =   shard(__lowerCAmelCase  )

								__magic_name__		:Tuple   =   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.045043945)  ) < 1E-3
												assert np.abs((np.abs(__lowerCAmelCase							,       dtype=np.floataa  ).sum() - 2347693.5)  ) < 5E-1







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Optional[int]   =   (
								    '''A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'''
								    ''' field, close up, split lighting, cinematic'''
								)

								__magic_name__		:str   =   jax.device_count()
								__magic_name__		:str   =   num_samples * [prompt]
								__magic_name__		:Dict   =   jax.random.split(jax.random.PRNGKey(0  )							,       __lowerCAmelCase  )

								__magic_name__  ,					__magic_name__		:Union[str, Any]   =   FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4'''							,       revision='''bf16'''							,       dtype=jnp.bfloataa							,       safety_checker=__lowerCAmelCase							,       )

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

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

								__magic_name__		:List[Any]   =   replicate(__lowerCAmelCase  )
								__magic_name__		:Any   =   pipeline.prepare_inputs(__lowerCAmelCase  )
								__magic_name__		:Any   =   shard(__lowerCAmelCase  )
								__magic_name__		:Any   =   pipeline(__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       jit=__lowerCAmelCase  ).images
								assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
								__magic_name__		:List[Any]   =   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_distilbert import DistilBertTokenizer


UpperCamelCase_    	=				logging.get_logger(__name__)

UpperCamelCase_    	=				{"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

UpperCamelCase_    	=				{
    "vocab_file": {
        "distilbert-base-uncased": "https://huggingface.co/distilbert-base-uncased/resolve/main/vocab.txt",
        "distilbert-base-uncased-distilled-squad": (
            "https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/vocab.txt"
        ),
        "distilbert-base-cased": "https://huggingface.co/distilbert-base-cased/resolve/main/vocab.txt",
        "distilbert-base-cased-distilled-squad": (
            "https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/vocab.txt"
        ),
        "distilbert-base-german-cased": "https://huggingface.co/distilbert-base-german-cased/resolve/main/vocab.txt",
        "distilbert-base-multilingual-cased": (
            "https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "distilbert-base-uncased": "https://huggingface.co/distilbert-base-uncased/resolve/main/tokenizer.json",
        "distilbert-base-uncased-distilled-squad": (
            "https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/tokenizer.json"
        ),
        "distilbert-base-cased": "https://huggingface.co/distilbert-base-cased/resolve/main/tokenizer.json",
        "distilbert-base-cased-distilled-squad": (
            "https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/tokenizer.json"
        ),
        "distilbert-base-german-cased": (
            "https://huggingface.co/distilbert-base-german-cased/resolve/main/tokenizer.json"
        ),
        "distilbert-base-multilingual-cased": (
            "https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/tokenizer.json"
        ),
    },
}

UpperCamelCase_    	=				{
    "distilbert-base-uncased": 5_1_2,
    "distilbert-base-uncased-distilled-squad": 5_1_2,
    "distilbert-base-cased": 5_1_2,
    "distilbert-base-cased-distilled-squad": 5_1_2,
    "distilbert-base-german-cased": 5_1_2,
    "distilbert-base-multilingual-cased": 5_1_2,
}


UpperCamelCase_    	=				{
    "distilbert-base-uncased": {"do_lower_case": True},
    "distilbert-base-uncased-distilled-squad": {"do_lower_case": True},
    "distilbert-base-cased": {"do_lower_case": False},
    "distilbert-base-cased-distilled-squad": {"do_lower_case": False},
    "distilbert-base-german-cased": {"do_lower_case": False},
    "distilbert-base-multilingual-cased": {"do_lower_case": False},
}






class 							_a				(			SCREAMING_SNAKE_CASE ):





       '''simple docstring'''






       A						:				List[Any]        =						VOCAB_FILES_NAMES
       A						:				Dict        =						PRETRAINED_VOCAB_FILES_MAP
       A						:				Optional[int]        =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       A						:				Optional[Any]        =						PRETRAINED_INIT_CONFIGURATION
       A						:				Optional[int]        =						['''input_ids''', '''attention_mask''']
       A						:				List[Any]        =						DistilBertTokenizer







       def __init__(      self,							A=None,							A=None,							A=True,							A="[UNK]",							A="[SEP]",							A="[PAD]",							A="[CLS]",							A="[MASK]",							A=True,							A=None,							**A,							):







             '''simple docstring'''







             super().__init__(
                 A,							tokenizer_file=A,							do_lower_case=A,							unk_token=A,							sep_token=A,							pad_token=A,							cls_token=A,							mask_token=A,							tokenize_chinese_chars=A,							strip_accents=A,							**A,							)

             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						json.loads(self.backend_tokenizer.normalizer.__getstate__()  )
             if (
                 normalizer_state.get('lowercase',							A  ) != do_lower_case
                 or normalizer_state.get('strip_accents',							A  ) != strip_accents
                 or normalizer_state.get('handle_chinese_chars',							A  ) != tokenize_chinese_chars
             ):
                   SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						getattr(A,							normalizer_state.pop('type'  )  )
                   SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						do_lower_case
                   SCREAMING_SNAKE_CASE   :      List[str]    					=						strip_accents
                   SCREAMING_SNAKE_CASE   :      List[str]    					=						tokenize_chinese_chars
                   SCREAMING_SNAKE_CASE   :      Dict    					=						normalizer_class(**A  )

             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						do_lower_case







       def        UpperCamelCase_							(      self,							A,							A=None  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      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        UpperCamelCase_							(      self,							A,							A = None  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      List[str]    					=						[self.sep_token_id]
             SCREAMING_SNAKE_CASE   :      str    					=						[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        UpperCamelCase_							(      self,							A,							A = None  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      List[str]    					=						self._tokenizer.model.save(A,							name=A  )
             return tuple(A  )


