





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



 def 				UpperCamelCase			(							self ):
        with tempfile.TemporaryDirectory() as tmpdirname:
               # pipeline has Flax weights
               A__											=						FlaxDiffusionPipeline.from_pretrained(
                   '''hf-internal-testing/tiny-stable-diffusion-pipe''',safety_checker=lowerCAmelCase__,cache_dir=lowerCAmelCase__ )

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



 def 				UpperCamelCase			(							self ):
        A__											=						FlaxStableDiffusionPipeline.from_pretrained(
            '''hf-internal-testing/tiny-stable-diffusion-pipe''',safety_checker=lowerCAmelCase__ )

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

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

        A__											=						jax.device_count()
        A__											=						num_samples * [prompt]
        A__											=						pipeline.prepare_inputs(lowerCAmelCase__ )

        # shard inputs and rng
        A__											=						replicate(lowerCAmelCase__ )
        A__											=						jax.random.split(lowerCAmelCase__,lowerCAmelCase__ )
        A__											=						shard(lowerCAmelCase__ )

        A__											=						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() - 4_9947.875 ) < 5E-1

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



 def 				UpperCamelCase			(							self ):
        A__											=						FlaxStableDiffusionPipeline.from_pretrained(
            '''CompVis/stable-diffusion-v1-4''',revision='''flax''',safety_checker=lowerCAmelCase__ )

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

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

        A__											=						jax.device_count()
        A__											=						num_samples * [prompt]
        A__											=						pipeline.prepare_inputs(lowerCAmelCase__ )

        # shard inputs and rng
        A__											=						replicate(lowerCAmelCase__ )
        A__											=						jax.random.split(lowerCAmelCase__,lowerCAmelCase__ )
        A__											=						shard(lowerCAmelCase__ )

        A__											=						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() - 238_3808.2) ) < 5E-1



 def 				UpperCamelCase			(							self ):
        A__											=						FlaxStableDiffusionPipeline.from_pretrained(
            '''CompVis/stable-diffusion-v1-4''',revision='''bf16''',dtype=jnp.bfloataa,safety_checker=lowerCAmelCase__ )

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

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

        A__											=						jax.device_count()
        A__											=						num_samples * [prompt]
        A__											=						pipeline.prepare_inputs(lowerCAmelCase__ )

        # shard inputs and rng
        A__											=						replicate(lowerCAmelCase__ )
        A__											=						jax.random.split(lowerCAmelCase__,lowerCAmelCase__ )
        A__											=						shard(lowerCAmelCase__ )

        A__											=						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() - 237_3516.75) ) < 5E-1



 def 				UpperCamelCase			(							self ):
        A__											=						FlaxStableDiffusionPipeline.from_pretrained(
            '''CompVis/stable-diffusion-v1-4''',revision='''bf16''',dtype=jnp.bfloataa )

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

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

        A__											=						jax.device_count()
        A__											=						num_samples * [prompt]
        A__											=						pipeline.prepare_inputs(lowerCAmelCase__ )

        # shard inputs and rng
        A__											=						replicate(lowerCAmelCase__ )
        A__											=						jax.random.split(lowerCAmelCase__,lowerCAmelCase__ )
        A__											=						shard(lowerCAmelCase__ )

        A__											=						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() - 237_3516.75) ) < 5E-1



 def 				UpperCamelCase			(							self ):
        A__											=						FlaxDDIMScheduler(
            beta_start=0.00085,beta_end=0.012,beta_schedule='''scaled_linear''',set_alpha_to_one=lowerCAmelCase__,steps_offset=1,)

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

        A__											=						scheduler_state

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

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

        A__											=						jax.device_count()
        A__											=						num_samples * [prompt]
        A__											=						pipeline.prepare_inputs(lowerCAmelCase__ )

        # shard inputs and rng
        A__											=						replicate(lowerCAmelCase__ )
        A__											=						jax.random.split(lowerCAmelCase__,lowerCAmelCase__ )
        A__											=						shard(lowerCAmelCase__ )

        A__											=						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() - 234_7693.5) ) < 5E-1




 def 				UpperCamelCase			(							self ):
        A__											=						(
            "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
            " field, close up, split lighting, cinematic"
        )

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

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

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

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

        A__											=						replicate(lowerCAmelCase__ )
        A__											=						pipeline.prepare_inputs(lowerCAmelCase__ )
        A__											=						shard(lowerCAmelCase__ )
        A__											=						pipeline(lowerCAmelCase__,lowerCAmelCase__,lowerCAmelCase__,jit=lowerCAmelCase__ ).images
        assert images_eff.shape == (num_samples, 1, 512, 512, 3)
        A__											=						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












"""simple docstring"""




from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_herbert import HerbertTokenizer


a						: List[str]     	=			logging.get_logger(__name__)

a						: List[Any]     	=			{'''vocab_file''': '''vocab.json''', '''merges_file''': '''merges.txt''', '''tokenizer_file''': '''tokenizer.json'''}

a						: str     	=			{
    '''vocab_file''': {
        '''allegro/herbert-base-cased''': '''https://huggingface.co/allegro/herbert-base-cased/resolve/main/vocab.json'''
    },
    '''merges_file''': {
        '''allegro/herbert-base-cased''': '''https://huggingface.co/allegro/herbert-base-cased/resolve/main/merges.txt'''
    },
}

a						: Tuple     	=			{'''allegro/herbert-base-cased''': 514}
a						: Optional[int]     	=			{}


class        __UpperCamelCase				(						a__						):
							lowerCamelCase   :					str      					=VOCAB_FILES_NAMES
							lowerCamelCase   :					Dict      					=PRETRAINED_VOCAB_FILES_MAP
							lowerCamelCase   :					Dict      					=PRETRAINED_INIT_CONFIGURATION
							lowerCamelCase   :					int      					=PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							lowerCamelCase   :					Optional[Any]      					=HerbertTokenizer






							def __init__(   self      ,					lowerCAmelCase__=None      ,					lowerCAmelCase__=None      ,					lowerCAmelCase__=None      ,					lowerCAmelCase__="<s>"      ,					lowerCAmelCase__="<unk>"      ,					lowerCAmelCase__="<pad>"      ,					lowerCAmelCase__="<mask>"      ,					lowerCAmelCase__="</s>"      ,					**lowerCAmelCase__      ,					)	-> Optional[int]:
								super().__init__(
								    lowerCAmelCase__      ,					lowerCAmelCase__      ,					tokenizer_file=lowerCAmelCase__      ,					cls_token=lowerCAmelCase__      ,					unk_token=lowerCAmelCase__      ,					pad_token=lowerCAmelCase__      ,					mask_token=lowerCAmelCase__      ,					sep_token=lowerCAmelCase__      ,					**lowerCAmelCase__      ,					)






							def     __a   (   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = None      )	-> List[int]:
								a      :		Optional[Any]	=				[self.cls_token_id]
								a      :		Any	=				[self.sep_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     __a   (   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = None      ,					lowerCAmelCase__ = False      )	-> List[int]:
								if already_has_special_tokens:
									return super().get_special_tokens_mask(
									    token_ids_a=lowerCAmelCase__      ,					token_ids_a=lowerCAmelCase__      ,					already_has_special_tokens=lowerCAmelCase__      )

								if token_ids_a is None:
									return [1] + ([0] * len(lowerCAmelCase__      )) + [1]
								return [1] + ([0] * len(lowerCAmelCase__      )) + [1] + ([0] * len(lowerCAmelCase__      )) + [1]






							def     __a   (   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = None      )	-> List[int]:
								a      :		Dict	=				[self.sep_token_id]
								a      :		Tuple	=				[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     __a   (   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = None      )	-> Tuple[str]:
								a      :		List[str]	=				self._tokenizer.model.save(lowerCAmelCase__      ,					name=lowerCAmelCase__      )
								return tuple(lowerCAmelCase__      )


