






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    _UpperCAmelCase      ( unittest.TestCase			):

				"""simple docstring"""




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

								'''simple docstring'''
								with tempfile.TemporaryDirectory() as tmpdirname:
												# pipeline has Flax weights
												lowercase__     				=    FlaxDiffusionPipeline.from_pretrained(
												    '''hf-internal-testing/tiny-stable-diffusion-pipe''',						safety_checker=lowerCamelCase,						cache_dir=lowerCamelCase						)

												lowercase__     				=    [t[-1] for t in os.walk(os.path.join(lowerCamelCase,						os.listdir(lowerCamelCase						)[0],						'''snapshots'''						)						)]
												lowercase__     				=    [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    _UpperCAmelCase      ( unittest.TestCase			):

				"""simple docstring"""




				def        lowercase__     (  self		:					int						):

								'''simple docstring'''
								lowercase__				,		lowercase__     				=    FlaxStableDiffusionPipeline.from_pretrained(
								    '''hf-internal-testing/tiny-stable-diffusion-pipe''',						safety_checker=lowerCamelCase						)

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

								lowercase__     				=    jax.random.PRNGKey(0						)
								lowercase__     				=    4

								lowercase__     				=    jax.device_count()
								lowercase__     				=    num_samples * [prompt]
								lowercase__     				=    pipeline.prepare_inputs(lowerCamelCase						)

								# shard inputs and rng
								lowercase__     				=    replicate(lowerCamelCase						)
								lowercase__     				=    jax.random.split(lowerCamelCase,						lowerCamelCase						)
								lowercase__     				=    shard(lowerCamelCase						)

								lowercase__     				=    pipeline(lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						jit=lowerCamelCase						).images

								assert images.shape == (num_samples, 1, 64, 64, 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

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




				def        lowercase__     (  self		:					Any						):

								'''simple docstring'''
								lowercase__				,		lowercase__     				=    FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4''',						revision='''flax''',						safety_checker=lowerCamelCase						)

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

								lowercase__     				=    jax.random.PRNGKey(0						)
								lowercase__     				=    50

								lowercase__     				=    jax.device_count()
								lowercase__     				=    num_samples * [prompt]
								lowercase__     				=    pipeline.prepare_inputs(lowerCamelCase						)

								# shard inputs and rng
								lowercase__     				=    replicate(lowerCamelCase						)
								lowercase__     				=    jax.random.split(lowerCamelCase,						lowerCamelCase						)
								lowercase__     				=    shard(lowerCamelCase						)

								lowercase__     				=    pipeline(lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						jit=lowerCamelCase						).images

								assert images.shape == (num_samples, 1, 512, 512, 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        lowercase__     (  self		:					int						):

								'''simple docstring'''
								lowercase__				,		lowercase__     				=    FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4''',						revision='''bf16''',						dtype=jnp.bfloataa,						safety_checker=lowerCamelCase						)

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

								lowercase__     				=    jax.random.PRNGKey(0						)
								lowercase__     				=    50

								lowercase__     				=    jax.device_count()
								lowercase__     				=    num_samples * [prompt]
								lowercase__     				=    pipeline.prepare_inputs(lowerCamelCase						)

								# shard inputs and rng
								lowercase__     				=    replicate(lowerCamelCase						)
								lowercase__     				=    jax.random.split(lowerCamelCase,						lowerCamelCase						)
								lowercase__     				=    shard(lowerCamelCase						)

								lowercase__     				=    pipeline(lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						jit=lowerCamelCase						).images

								assert images.shape == (num_samples, 1, 512, 512, 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        lowercase__     (  self		:					Dict						):

								'''simple docstring'''
								lowercase__				,		lowercase__     				=    FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4''',						revision='''bf16''',						dtype=jnp.bfloataa						)

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

								lowercase__     				=    jax.random.PRNGKey(0						)
								lowercase__     				=    50

								lowercase__     				=    jax.device_count()
								lowercase__     				=    num_samples * [prompt]
								lowercase__     				=    pipeline.prepare_inputs(lowerCamelCase						)

								# shard inputs and rng
								lowercase__     				=    replicate(lowerCamelCase						)
								lowercase__     				=    jax.random.split(lowerCamelCase,						lowerCamelCase						)
								lowercase__     				=    shard(lowerCamelCase						)

								lowercase__     				=    pipeline(lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						jit=lowerCamelCase						).images

								assert images.shape == (num_samples, 1, 512, 512, 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        lowercase__     (  self		:					Optional[int]						):

								'''simple docstring'''
								lowercase__     				=    FlaxDDIMScheduler(
								    beta_start=0.00085,						beta_end=0.012,						beta_schedule='''scaled_linear''',						set_alpha_to_one=lowerCamelCase,						steps_offset=1,						)

								lowercase__				,		lowercase__     				=    FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4''',						revision='''bf16''',						dtype=jnp.bfloataa,						scheduler=lowerCamelCase,						safety_checker=lowerCamelCase,						)
								lowercase__     				=    scheduler.create_state()

								lowercase__     				=    scheduler_state

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

								lowercase__     				=    jax.random.PRNGKey(0						)
								lowercase__     				=    50

								lowercase__     				=    jax.device_count()
								lowercase__     				=    num_samples * [prompt]
								lowercase__     				=    pipeline.prepare_inputs(lowerCamelCase						)

								# shard inputs and rng
								lowercase__     				=    replicate(lowerCamelCase						)
								lowercase__     				=    jax.random.split(lowerCamelCase,						lowerCamelCase						)
								lowercase__     				=    shard(lowerCamelCase						)

								lowercase__     				=    pipeline(lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						jit=lowerCamelCase						).images

								assert images.shape == (num_samples, 1, 512, 512, 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        lowercase__     (  self		:					str						):

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

								lowercase__     				=    jax.device_count()
								lowercase__     				=    num_samples * [prompt]
								lowercase__     				=    jax.random.split(jax.random.PRNGKey(0						),						lowerCamelCase						)

								lowercase__				,		lowercase__     				=    FlaxStableDiffusionPipeline.from_pretrained(
								    '''CompVis/stable-diffusion-v1-4''',						revision='''bf16''',						dtype=jnp.bfloataa,						safety_checker=lowerCamelCase,						)

								lowercase__     				=    replicate(lowerCamelCase						)
								lowercase__     				=    pipeline.prepare_inputs(lowerCamelCase						)
								lowercase__     				=    shard(lowerCamelCase						)
								lowercase__     				=    pipeline(lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						jit=lowerCamelCase						).images
								assert images.shape == (num_samples, 1, 512, 512, 3)
								lowercase__     				=    images[2, 0, 256, 10:17, 1]

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

								lowercase__     				=    replicate(lowerCamelCase						)
								lowercase__     				=    pipeline.prepare_inputs(lowerCamelCase						)
								lowercase__     				=    shard(lowerCamelCase						)
								lowercase__     				=    pipeline(lowerCamelCase,						lowerCamelCase,						lowerCamelCase,						jit=lowerCamelCase						).images
								assert images_eff.shape == (num_samples, 1, 512, 512, 3)
								lowercase__     				=    images[2, 0, 256, 10:17, 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












import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_bert import BertTokenizer


A__			:    Dict         =				logging.get_logger(__name__)

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

A__			:    Optional[int]         =				{
    'vocab_file': {
        'bert-base-uncased': 'https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt',
        'bert-large-uncased': 'https://huggingface.co/bert-large-uncased/resolve/main/vocab.txt',
        'bert-base-cased': 'https://huggingface.co/bert-base-cased/resolve/main/vocab.txt',
        'bert-large-cased': 'https://huggingface.co/bert-large-cased/resolve/main/vocab.txt',
        'bert-base-multilingual-uncased': (
            'https://huggingface.co/bert-base-multilingual-uncased/resolve/main/vocab.txt'
        ),
        'bert-base-multilingual-cased': 'https://huggingface.co/bert-base-multilingual-cased/resolve/main/vocab.txt',
        'bert-base-chinese': 'https://huggingface.co/bert-base-chinese/resolve/main/vocab.txt',
        'bert-base-german-cased': 'https://huggingface.co/bert-base-german-cased/resolve/main/vocab.txt',
        'bert-large-uncased-whole-word-masking': (
            'https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/vocab.txt'
        ),
        'bert-large-cased-whole-word-masking': (
            'https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/vocab.txt'
        ),
        'bert-large-uncased-whole-word-masking-finetuned-squad': (
            'https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt'
        ),
        'bert-large-cased-whole-word-masking-finetuned-squad': (
            'https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt'
        ),
        'bert-base-cased-finetuned-mrpc': (
            'https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/vocab.txt'
        ),
        'bert-base-german-dbmdz-cased': 'https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/vocab.txt',
        'bert-base-german-dbmdz-uncased': (
            'https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/vocab.txt'
        ),
        'TurkuNLP/bert-base-finnish-cased-v1': (
            'https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/vocab.txt'
        ),
        'TurkuNLP/bert-base-finnish-uncased-v1': (
            'https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/vocab.txt'
        ),
        'wietsedv/bert-base-dutch-cased': (
            'https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/vocab.txt'
        ),
    },
    'tokenizer_file': {
        'bert-base-uncased': 'https://huggingface.co/bert-base-uncased/resolve/main/tokenizer.json',
        'bert-large-uncased': 'https://huggingface.co/bert-large-uncased/resolve/main/tokenizer.json',
        'bert-base-cased': 'https://huggingface.co/bert-base-cased/resolve/main/tokenizer.json',
        'bert-large-cased': 'https://huggingface.co/bert-large-cased/resolve/main/tokenizer.json',
        'bert-base-multilingual-uncased': (
            'https://huggingface.co/bert-base-multilingual-uncased/resolve/main/tokenizer.json'
        ),
        'bert-base-multilingual-cased': (
            'https://huggingface.co/bert-base-multilingual-cased/resolve/main/tokenizer.json'
        ),
        'bert-base-chinese': 'https://huggingface.co/bert-base-chinese/resolve/main/tokenizer.json',
        'bert-base-german-cased': 'https://huggingface.co/bert-base-german-cased/resolve/main/tokenizer.json',
        'bert-large-uncased-whole-word-masking': (
            'https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/tokenizer.json'
        ),
        'bert-large-cased-whole-word-masking': (
            'https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/tokenizer.json'
        ),
        'bert-large-uncased-whole-word-masking-finetuned-squad': (
            'https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/tokenizer.json'
        ),
        'bert-large-cased-whole-word-masking-finetuned-squad': (
            'https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/tokenizer.json'
        ),
        'bert-base-cased-finetuned-mrpc': (
            'https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/tokenizer.json'
        ),
        'bert-base-german-dbmdz-cased': (
            'https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/tokenizer.json'
        ),
        'bert-base-german-dbmdz-uncased': (
            'https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/tokenizer.json'
        ),
        'TurkuNLP/bert-base-finnish-cased-v1': (
            'https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/tokenizer.json'
        ),
        'TurkuNLP/bert-base-finnish-uncased-v1': (
            'https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/tokenizer.json'
        ),
        'wietsedv/bert-base-dutch-cased': (
            'https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/tokenizer.json'
        ),
    },
}

A__			:    List[str]         =				{
    'bert-base-uncased': 5_12,
    'bert-large-uncased': 5_12,
    'bert-base-cased': 5_12,
    'bert-large-cased': 5_12,
    'bert-base-multilingual-uncased': 5_12,
    'bert-base-multilingual-cased': 5_12,
    'bert-base-chinese': 5_12,
    'bert-base-german-cased': 5_12,
    'bert-large-uncased-whole-word-masking': 5_12,
    'bert-large-cased-whole-word-masking': 5_12,
    'bert-large-uncased-whole-word-masking-finetuned-squad': 5_12,
    'bert-large-cased-whole-word-masking-finetuned-squad': 5_12,
    'bert-base-cased-finetuned-mrpc': 5_12,
    'bert-base-german-dbmdz-cased': 5_12,
    'bert-base-german-dbmdz-uncased': 5_12,
    'TurkuNLP/bert-base-finnish-cased-v1': 5_12,
    'TurkuNLP/bert-base-finnish-uncased-v1': 5_12,
    'wietsedv/bert-base-dutch-cased': 5_12,
}

A__			:    Optional[int]         =				{
    'bert-base-uncased': {'do_lower_case': True},
    'bert-large-uncased': {'do_lower_case': True},
    'bert-base-cased': {'do_lower_case': False},
    'bert-large-cased': {'do_lower_case': False},
    'bert-base-multilingual-uncased': {'do_lower_case': True},
    'bert-base-multilingual-cased': {'do_lower_case': False},
    'bert-base-chinese': {'do_lower_case': False},
    'bert-base-german-cased': {'do_lower_case': False},
    'bert-large-uncased-whole-word-masking': {'do_lower_case': True},
    'bert-large-cased-whole-word-masking': {'do_lower_case': False},
    'bert-large-uncased-whole-word-masking-finetuned-squad': {'do_lower_case': True},
    'bert-large-cased-whole-word-masking-finetuned-squad': {'do_lower_case': False},
    'bert-base-cased-finetuned-mrpc': {'do_lower_case': False},
    'bert-base-german-dbmdz-cased': {'do_lower_case': False},
    'bert-base-german-dbmdz-uncased': {'do_lower_case': True},
    'TurkuNLP/bert-base-finnish-cased-v1': {'do_lower_case': False},
    'TurkuNLP/bert-base-finnish-uncased-v1': {'do_lower_case': True},
    'wietsedv/bert-base-dutch-cased': {'do_lower_case': False},
}
class    _UpperCAmelCase      ( A__			):

				"""simple docstring"""




				lowercase__       							= VOCAB_FILES_NAMES
				lowercase__       							= PRETRAINED_VOCAB_FILES_MAP
				lowercase__       							= PRETRAINED_INIT_CONFIGURATION
				lowercase__       							= PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				lowercase__       							= BertTokenizer




				def __init__(  self		:					Any,						lowerCamelCase		:					Optional[Any]=None,						lowerCamelCase		:					Any=None,						lowerCamelCase		:					Tuple=True,						lowerCamelCase		:					Dict="[UNK]",						lowerCamelCase		:					Any="[SEP]",						lowerCamelCase		:					List[Any]="[PAD]",						lowerCamelCase		:					Optional[Any]="[CLS]",						lowerCamelCase		:					Dict="[MASK]",						lowerCamelCase		:					List[Any]=True,						lowerCamelCase		:					Tuple=None,						**lowerCamelCase		:					Dict,						):

								'''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,						tokenize_chinese_chars=lowerCamelCase,						strip_accents=lowerCamelCase,						**lowerCamelCase,						)

								lowercase__     				=    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
								):
												lowercase__     				=    getattr(lowerCamelCase,						normalizer_state.pop('''type'''						)						)
												lowercase__     				=    do_lower_case
												lowercase__     				=    strip_accents
												lowercase__     				=    tokenize_chinese_chars
												lowercase__     				=    normalizer_class(**lowerCamelCase						)

								lowercase__     				=    do_lower_case




				def        lowercase__     (  self		:					Any,						lowerCamelCase		:					List[Any],						lowerCamelCase		:					Dict=None						):

								'''simple docstring'''
								lowercase__     				=    [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        lowercase__     (  self		:					List[Any],						lowerCamelCase		:					List[int],						lowerCamelCase		:					Optional[List[int]] = None						):

								'''simple docstring'''
								lowercase__     				=    [self.sep_token_id]
								lowercase__     				=    [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        lowercase__     (  self		:					Any,						lowerCamelCase		:					str,						lowerCamelCase		:					Optional[str] = None						):

								'''simple docstring'''
								lowercase__     				=    self._tokenizer.model.save(lowerCamelCase,						name=lowerCamelCase						)
								return tuple(lowerCamelCase						)




