"""simple docstring"""







import json
import os
import shutil
import tempfile
import unittest

import numpy as np

from transformers import BertTokenizerFast
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES, BertTokenizer
from transformers.testing_utils import require_tokenizers, require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
       from PIL import Image

       from transformers import VisionTextDualEncoderProcessor, ViTImageProcessor







@require_tokenizers
@require_vision
class  UpperCamelCase	(    unittest.TestCase						):
   def     _UpperCAmelCase     (    self    )			->   Optional[int]:


        '''simple docstring'''


        lowercase_			:       List[str]									=			tempfile.mkdtemp()

        # fmt: off
        lowercase_			:       Dict									=			['[UNK]', '[CLS]', '[SEP]', '[PAD]', '[MASK]', 'want', '##want', '##ed', 'wa', 'un', 'runn', '##ing', ',', 'low', 'lowest']
        # fmt: on
        lowercase_			:       Dict									=			os.path.join(self.tmpdirname  ,VOCAB_FILES_NAMES['vocab_file']    )
        with open(self.vocab_file  ,'w'  ,encoding='utf-8'    ) as vocab_writer:
             vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]    )    )

        lowercase_			:       Optional[int]									=			{
            'do_resize': True,
            'size': {'height': 18, 'width': 18},
            'do_normalize': True,
            'image_mean': [0.5, 0.5, 0.5],
            'image_std': [0.5, 0.5, 0.5],
        }
        lowercase_			:       List[Any]									=			os.path.join(self.tmpdirname  ,__UpperCamelCase    )
        with open(self.image_processor_file  ,'w'  ,encoding='utf-8'    ) as fp:
             json.dump(__UpperCamelCase  ,__UpperCamelCase    )
   def     _UpperCAmelCase     (    self  ,**__UpperCamelCase    )			->   Optional[int]:


        '''simple docstring'''


        return BertTokenizer.from_pretrained(self.tmpdirname  ,**__UpperCamelCase    )
   def     _UpperCAmelCase     (    self  ,**__UpperCamelCase    )			->   List[Any]:


        '''simple docstring'''


        return ViTImageProcessor.from_pretrained(self.tmpdirname  ,**__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   Optional[Any]:


        '''simple docstring'''


        shutil.rmtree(self.tmpdirname    )
   def     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


        lowercase_			:       str									=			[np.random.randint(255  ,size=(3, 30, 400)  ,dtype=np.uinta    )]

        lowercase_			:       int									=			[Image.fromarray(np.moveaxis(__UpperCamelCase  ,0  ,-1    )    ) for x in image_inputs]

        return image_inputs
   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        lowercase_			:       List[Any]									=			self.get_tokenizer()
        lowercase_			:       int									=			self.get_image_processor()

        lowercase_			:       Optional[Any]									=			VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase  ,image_processor=__UpperCamelCase    )

        processor.save_pretrained(self.tmpdirname    )
        lowercase_			:       Optional[Any]									=			VisionTextDualEncoderProcessor.from_pretrained(self.tmpdirname    )

        self.assertEqual(processor.tokenizer.get_vocab()  ,tokenizer.get_vocab()    )
        self.assertIsInstance(processor.tokenizer  ,(BertTokenizer, BertTokenizerFast)    )

        self.assertEqual(processor.image_processor.to_json_string()  ,image_processor.to_json_string()    )
        self.assertIsInstance(processor.image_processor  ,__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			VisionTextDualEncoderProcessor(
            tokenizer=self.get_tokenizer()  ,image_processor=self.get_image_processor()    )
        processor.save_pretrained(self.tmpdirname    )

        lowercase_			:       Dict									=			self.get_tokenizer(bos_token='(BOS)'  ,eos_token='(EOS)'    )
        lowercase_			:       Union[str, Any]									=			self.get_image_processor(do_normalize=__UpperCamelCase  ,padding_value=1.0    )

        lowercase_			:       int									=			VisionTextDualEncoderProcessor.from_pretrained(
            self.tmpdirname  ,bos_token='(BOS)'  ,eos_token='(EOS)'  ,do_normalize=__UpperCamelCase  ,padding_value=1.0    )

        self.assertEqual(processor.tokenizer.get_vocab()  ,tokenizer_add_kwargs.get_vocab()    )
        self.assertIsInstance(processor.tokenizer  ,(BertTokenizer, BertTokenizerFast)    )

        self.assertEqual(processor.image_processor.to_json_string()  ,image_processor_add_kwargs.to_json_string()    )
        self.assertIsInstance(processor.image_processor  ,__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			self.get_image_processor()
        lowercase_			:       int									=			self.get_tokenizer()

        lowercase_			:       List[Any]									=			VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase  ,image_processor=__UpperCamelCase    )

        lowercase_			:       Union[str, Any]									=			self.prepare_image_inputs()

        lowercase_			:       Optional[Any]									=			image_processor(__UpperCamelCase  ,return_tensors='np'    )
        lowercase_			:       Tuple									=			processor(images=__UpperCamelCase  ,return_tensors='np'    )

        for key in input_feat_extract.keys():
             self.assertAlmostEqual(input_feat_extract[key].sum()  ,input_processor[key].sum()  ,delta=1e-2    )
   def     _UpperCAmelCase     (    self    )			->   Optional[int]:


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			self.get_image_processor()
        lowercase_			:       Tuple									=			self.get_tokenizer()

        lowercase_			:       List[Any]									=			VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase  ,image_processor=__UpperCamelCase    )

        lowercase_			:       Any									=			'lower newer'

        lowercase_			:       str									=			processor(text=__UpperCamelCase    )

        lowercase_			:       Dict									=			tokenizer(__UpperCamelCase    )

        for key in encoded_tok.keys():
             self.assertListEqual(encoded_tok[key]  ,encoded_processor[key]    )
   def     _UpperCAmelCase     (    self    )			->   Tuple:


        '''simple docstring'''


        lowercase_			:       int									=			self.get_image_processor()
        lowercase_			:       List[str]									=			self.get_tokenizer()

        lowercase_			:       str									=			VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase  ,image_processor=__UpperCamelCase    )

        lowercase_			:       Union[str, Any]									=			'lower newer'
        lowercase_			:       str									=			self.prepare_image_inputs()

        lowercase_			:       int									=			processor(text=__UpperCamelCase  ,images=__UpperCamelCase    )

        self.assertListEqual(list(inputs.keys()    )  ,['input_ids', 'token_type_ids', 'attention_mask', 'pixel_values']    )

        # test if it raises when no input is passed
        with self.assertRaises(__UpperCamelCase    ):
             processor()
   def     _UpperCAmelCase     (    self    )			->   Optional[int]:


        '''simple docstring'''


        lowercase_			:       int									=			self.get_image_processor()
        lowercase_			:       List[Any]									=			self.get_tokenizer()

        lowercase_			:       Optional[Any]									=			VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase  ,image_processor=__UpperCamelCase    )

        lowercase_			:       str									=			[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

        lowercase_			:       Dict									=			processor.batch_decode(__UpperCamelCase    )
        lowercase_			:       Optional[int]									=			tokenizer.batch_decode(__UpperCamelCase    )

        self.assertListEqual(__UpperCamelCase  ,__UpperCamelCase    )







   def     _UpperCAmelCase     (    self    )			->   Any:


        '''simple docstring'''


        lowercase_			:       Dict									=			self.get_image_processor()
        lowercase_			:       Tuple									=			self.get_tokenizer()

        lowercase_			:       Tuple									=			VisionTextDualEncoderProcessor(tokenizer=__UpperCamelCase  ,image_processor=__UpperCamelCase    )

        lowercase_			:       int									=			'lower newer'
        lowercase_			:       Optional[int]									=			self.prepare_image_inputs()

        lowercase_			:       int									=			processor(text=__UpperCamelCase  ,images=__UpperCamelCase    )

        self.assertListEqual(list(inputs.keys()    )  ,processor.model_input_names    )


"""simple docstring"""







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						):
   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


        with tempfile.TemporaryDirectory() as tmpdirname:
             # pipeline has Flax weights
             lowercase_			:       Optional[int]									=			FlaxDiffusionPipeline.from_pretrained(
                 'hf-internal-testing/tiny-stable-diffusion-pipe'  ,safety_checker=__UpperCamelCase  ,cache_dir=__UpperCamelCase    )

             lowercase_			:       Dict									=			[t[-1] for t in os.walk(os.path.join(__UpperCamelCase  ,os.listdir(__UpperCamelCase    )[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  UpperCamelCase	(    unittest.TestCase						):
   def     _UpperCAmelCase     (    self    )			->   Optional[int]:


        '''simple docstring'''


        lowercase_							,						lowercase_			:       List[Any]									=			FlaxStableDiffusionPipeline.from_pretrained(
            'hf-internal-testing/tiny-stable-diffusion-pipe'  ,safety_checker=__UpperCamelCase    )

        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_			:       Optional[Any]									=			jax.random.PRNGKey(0    )
        lowercase_			:       Optional[int]									=			4

        lowercase_			:       int									=			jax.device_count()
        lowercase_			:       Optional[int]									=			num_samples * [prompt]
        lowercase_			:       List[Any]									=			pipeline.prepare_inputs(__UpperCamelCase    )

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

        lowercase_			:       List[Any]									=			pipeline(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,jit=__UpperCamelCase    ).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.151_4745    ) < 1e-3
             assert np.abs(np.abs(__UpperCamelCase  ,dtype=np.floataa    ).sum() - 4_9947.875    ) < 5e-1

        lowercase_			:       List[Any]									=			pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]    )    )    )
        assert len(__UpperCamelCase    ) == num_samples
   def     _UpperCAmelCase     (    self    )			->   Optional[Any]:


        '''simple docstring'''


        lowercase_							,						lowercase_			:       Any									=			FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'  ,revision='flax'  ,safety_checker=__UpperCamelCase    )

        lowercase_			:       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_			:       List[str]									=			jax.random.PRNGKey(0    )
        lowercase_			:       Tuple									=			50

        lowercase_			:       Optional[Any]									=			jax.device_count()
        lowercase_			:       List[Any]									=			num_samples * [prompt]
        lowercase_			:       Union[str, Any]									=			pipeline.prepare_inputs(__UpperCamelCase    )

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

        lowercase_			:       Optional[int]									=			pipeline(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,jit=__UpperCamelCase    ).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.0565_2401)    ) < 1e-3
             assert np.abs((np.abs(__UpperCamelCase  ,dtype=np.floataa    ).sum() - 238_3808.2)    ) < 5e-1
   def     _UpperCAmelCase     (    self    )			->   Union[str, Any]:


        '''simple docstring'''


        lowercase_							,						lowercase_			:       Tuple									=			FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'  ,revision='bf16'  ,dtype=jnp.bfloataa  ,safety_checker=__UpperCamelCase    )

        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_			:       Dict									=			jax.random.PRNGKey(0    )
        lowercase_			:       Optional[int]									=			50

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

        # shard inputs and rng
        lowercase_			:       int									=			replicate(__UpperCamelCase    )
        lowercase_			:       int									=			jax.random.split(__UpperCamelCase  ,__UpperCamelCase    )
        lowercase_			:       Union[str, Any]									=			shard(__UpperCamelCase    )

        lowercase_			:       List[str]									=			pipeline(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,jit=__UpperCamelCase    ).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.0400_3906)    ) < 1e-3
             assert np.abs((np.abs(__UpperCamelCase  ,dtype=np.floataa    ).sum() - 237_3516.75)    ) < 5e-1
   def     _UpperCAmelCase     (    self    )			->   Union[str, Any]:


        '''simple docstring'''


        lowercase_							,						lowercase_			:       str									=			FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'  ,revision='bf16'  ,dtype=jnp.bfloataa    )

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

        lowercase_			:       Optional[Any]									=			jax.random.PRNGKey(0    )
        lowercase_			:       Union[str, Any]									=			50

        lowercase_			:       Optional[Any]									=			jax.device_count()
        lowercase_			:       Tuple									=			num_samples * [prompt]
        lowercase_			:       Union[str, Any]									=			pipeline.prepare_inputs(__UpperCamelCase    )

        # shard inputs and rng
        lowercase_			:       List[Any]									=			replicate(__UpperCamelCase    )
        lowercase_			:       Union[str, Any]									=			jax.random.split(__UpperCamelCase  ,__UpperCamelCase    )
        lowercase_			:       Optional[int]									=			shard(__UpperCamelCase    )

        lowercase_			:       Optional[Any]									=			pipeline(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,jit=__UpperCamelCase    ).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.0400_3906)    ) < 1e-3
             assert np.abs((np.abs(__UpperCamelCase  ,dtype=np.floataa    ).sum() - 237_3516.75)    ) < 5e-1
   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


        lowercase_			:       Tuple									=			FlaxDDIMScheduler(
            beta_start=0.0_0085  ,beta_end=0.012  ,beta_schedule='scaled_linear'  ,set_alpha_to_one=__UpperCamelCase  ,steps_offset=1  ,)

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

        lowercase_			:       List[Any]									=			scheduler_state

        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_			:       str									=			jax.random.PRNGKey(0    )
        lowercase_			:       str									=			50

        lowercase_			:       List[Any]									=			jax.device_count()
        lowercase_			:       List[str]									=			num_samples * [prompt]
        lowercase_			:       Union[str, Any]									=			pipeline.prepare_inputs(__UpperCamelCase    )

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

        lowercase_			:       str									=			pipeline(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,jit=__UpperCamelCase    ).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.0_4504_3945)    ) < 1e-3
             assert np.abs((np.abs(__UpperCamelCase  ,dtype=np.floataa    ).sum() - 234_7693.5)    ) < 5e-1







   def     _UpperCAmelCase     (    self    )			->   Optional[Any]:


        '''simple docstring'''


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

        lowercase_			:       Union[str, Any]									=			jax.device_count()
        lowercase_			:       List[str]									=			num_samples * [prompt]
        lowercase_			:       int									=			jax.random.split(jax.random.PRNGKey(0    )  ,__UpperCamelCase    )

        lowercase_							,						lowercase_			:       Dict									=			FlaxStableDiffusionPipeline.from_pretrained(
            'CompVis/stable-diffusion-v1-4'  ,revision='bf16'  ,dtype=jnp.bfloataa  ,safety_checker=__UpperCamelCase  ,)

        lowercase_			:       str									=			replicate(__UpperCamelCase    )
        lowercase_			:       int									=			pipeline.prepare_inputs(__UpperCamelCase    )
        lowercase_			:       Optional[int]									=			shard(__UpperCamelCase    )
        lowercase_			:       Tuple									=			pipeline(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,jit=__UpperCamelCase    ).images
        assert images.shape == (num_samples, 1, 512, 512, 3)
        lowercase_			:       Union[str, Any]									=			images[2, 0, 256, 10:17, 1]

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

        lowercase_			:       List[Any]									=			replicate(__UpperCamelCase    )
        lowercase_			:       Union[str, Any]									=			pipeline.prepare_inputs(__UpperCamelCase    )
        lowercase_			:       List[Any]									=			shard(__UpperCamelCase    )
        lowercase_			:       Any									=			pipeline(__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,jit=__UpperCamelCase    ).images
        assert images_eff.shape == (num_samples, 1, 512, 512, 3)
        lowercase_			:       Optional[int]									=			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

