


'''simple docstring'''






import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import CLIPTokenizer, CLIPTokenizerFast
from transformers.models.clip.tokenization_clip import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
  from PIL import Image

  from transformers import OwlViTImageProcessor, OwlViTProcessor






@require_vision
class    __snake_case       (   unittest.TestCase):







    """simple docstring"""



    def __lowercase   (   self    :  Optional[Any]		)   ->  Any:
      lowerCAmelCase_			:      List[str]			=   tempfile.mkdtemp()

      # fmt: off
      lowerCAmelCase_			:      Optional[Any]			=   ["""""", """l""", """o""", """w""", """e""", """r""", """s""", """t""", """i""", """d""", """n""", """lo""", """l</w>""", """w</w>""", """r</w>""", """t</w>""", """low</w>""", """er</w>""", """lowest</w>""", """newer</w>""", """wider""", """<unk>""", """<|startoftext|>""", """<|endoftext|>"""]
      # fmt: on
      lowerCAmelCase_			:      Tuple			=   dict(zip(lowerCamelCase    ,						range(len(lowerCamelCase		)		)		)		)
      lowerCAmelCase_			:      List[str]			=   ["""#version: 0.2""", """l o""", """lo w</w>""", """e r</w>""", """"""]
      lowerCAmelCase_			:      str			=   {"""unk_token""": """<unk>"""}

      lowerCAmelCase_			:      Optional[int]			=   os.path.join(self.tmpdirname    ,						VOCAB_FILES_NAMES["""vocab_file"""]		)
      lowerCAmelCase_			:      Optional[int]			=   os.path.join(self.tmpdirname    ,						VOCAB_FILES_NAMES["""merges_file"""]		)
      with open(self.vocab_file    ,						"""w"""    ,						encoding="""utf-8"""		) as fp:
        fp.write(json.dumps(lowerCamelCase		) + """\n"""		)
      with open(self.merges_file    ,						"""w"""    ,						encoding="""utf-8"""		) as fp:
        fp.write("""\n""".join(lowerCamelCase		)		)

      lowerCAmelCase_			:      Tuple			=   {
          """do_resize""": True,
          """size""": 20,
          """do_center_crop""": True,
          """crop_size""": 18,
          """do_normalize""": True,
          """image_mean""": [0.48_145_466, 0.4_578_275, 0.40_821_073],
          """image_std""": [0.26_862_954, 0.26_130_258, 0.27_577_711],
      }
      lowerCAmelCase_			:      Dict			=   os.path.join(self.tmpdirname    ,						lowerCamelCase		)
      with open(self.image_processor_file    ,						"""w"""    ,						encoding="""utf-8"""		) as fp:
        json.dump(lowerCamelCase    ,						lowerCamelCase		)

    def __lowercase   (   self    :  Optional[Any]    ,						**lowerCamelCase    :  Any		)   ->  Optional[int]:
      return CLIPTokenizer.from_pretrained(self.tmpdirname    ,						pad_token="""!"""    ,						**lowerCamelCase		)

    def __lowercase   (   self    :  Optional[Any]    ,						**lowerCamelCase    :  List[Any]		)   ->  int:
      return CLIPTokenizerFast.from_pretrained(self.tmpdirname    ,						pad_token="""!"""    ,						**lowerCamelCase		)

    def __lowercase   (   self    :  Union[str, Any]    ,						**lowerCamelCase    :  List[Any]		)   ->  int:
      return OwlViTImageProcessor.from_pretrained(self.tmpdirname    ,						**lowerCamelCase		)

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

    def __lowercase   (   self    :  str		)   ->  List[Any]:
      lowerCAmelCase_			:      Dict			=   [np.random.randint(2_55    ,						size=(3, 30, 4_00)    ,						dtype=np.uinta		)]

      lowerCAmelCase_			:      Any			=   [Image.fromarray(np.moveaxis(lowerCamelCase    ,						0    ,						-1		)		) for x in image_inputs]

      return image_inputs

    def __lowercase   (   self    :  Tuple		)   ->  Dict:
      lowerCAmelCase_			:      Any			=   self.get_tokenizer()
      lowerCAmelCase_			:      Optional[Any]			=   self.get_rust_tokenizer()
      lowerCAmelCase_			:      Any			=   self.get_image_processor()

      lowerCAmelCase_			:      Union[str, Any]			=   OwlViTProcessor(tokenizer=lowerCamelCase    ,						image_processor=lowerCamelCase		)
      processor_slow.save_pretrained(self.tmpdirname		)
      lowerCAmelCase_			:      Optional[Any]			=   OwlViTProcessor.from_pretrained(self.tmpdirname    ,						use_fast=lowerCamelCase		)

      lowerCAmelCase_			:      Any			=   OwlViTProcessor(tokenizer=lowerCamelCase    ,						image_processor=lowerCamelCase		)
      processor_fast.save_pretrained(self.tmpdirname		)
      lowerCAmelCase_			:      int			=   OwlViTProcessor.from_pretrained(self.tmpdirname		)

      self.assertEqual(processor_slow.tokenizer.get_vocab()    ,						tokenizer_slow.get_vocab()		)
      self.assertEqual(processor_fast.tokenizer.get_vocab()    ,						tokenizer_fast.get_vocab()		)
      self.assertEqual(tokenizer_slow.get_vocab()    ,						tokenizer_fast.get_vocab()		)
      self.assertIsInstance(processor_slow.tokenizer    ,						lowerCamelCase		)
      self.assertIsInstance(processor_fast.tokenizer    ,						lowerCamelCase		)

      self.assertEqual(processor_slow.image_processor.to_json_string()    ,						image_processor.to_json_string()		)
      self.assertEqual(processor_fast.image_processor.to_json_string()    ,						image_processor.to_json_string()		)
      self.assertIsInstance(processor_slow.image_processor    ,						lowerCamelCase		)
      self.assertIsInstance(processor_fast.image_processor    ,						lowerCamelCase		)

    def __lowercase   (   self    :  str		)   ->  List[Any]:
      lowerCAmelCase_			:      Union[str, Any]			=   OwlViTProcessor(tokenizer=self.get_tokenizer()    ,						image_processor=self.get_image_processor()		)
      processor.save_pretrained(self.tmpdirname		)

      lowerCAmelCase_			:      str			=   self.get_tokenizer(bos_token="""(BOS)"""    ,						eos_token="""(EOS)"""		)
      lowerCAmelCase_			:      Any			=   self.get_image_processor(do_normalize=lowerCamelCase		)

      lowerCAmelCase_			:      Tuple			=   OwlViTProcessor.from_pretrained(
          self.tmpdirname    ,						bos_token="""(BOS)"""    ,						eos_token="""(EOS)"""    ,						do_normalize=lowerCamelCase		)

      self.assertEqual(processor.tokenizer.get_vocab()    ,						tokenizer_add_kwargs.get_vocab()		)
      self.assertIsInstance(processor.tokenizer    ,						lowerCamelCase		)

      self.assertEqual(processor.image_processor.to_json_string()    ,						image_processor_add_kwargs.to_json_string()		)
      self.assertIsInstance(processor.image_processor    ,						lowerCamelCase		)

    def __lowercase   (   self    :  int		)   ->  List[Any]:
      lowerCAmelCase_			:      Optional[int]			=   self.get_image_processor()
      lowerCAmelCase_			:      int			=   self.get_tokenizer()

      lowerCAmelCase_			:      Optional[Any]			=   OwlViTProcessor(tokenizer=lowerCamelCase    ,						image_processor=lowerCamelCase		)

      lowerCAmelCase_			:      int			=   self.prepare_image_inputs()

      lowerCAmelCase_			:      str			=   image_processor(lowerCamelCase    ,						return_tensors="""np"""		)
      lowerCAmelCase_			:      Any			=   processor(images=lowerCamelCase    ,						return_tensors="""np"""		)

      for key in input_image_proc.keys():
        self.assertAlmostEqual(input_image_proc[key].sum()    ,						input_processor[key].sum()    ,						delta=1E-2		)

    def __lowercase   (   self    :  Any		)   ->  List[str]:
      lowerCAmelCase_			:      List[Any]			=   self.get_image_processor()
      lowerCAmelCase_			:      List[str]			=   self.get_tokenizer()

      lowerCAmelCase_			:      List[str]			=   OwlViTProcessor(tokenizer=lowerCamelCase    ,						image_processor=lowerCamelCase		)

      lowerCAmelCase_			:      Tuple			=   """lower newer"""

      lowerCAmelCase_			:      Any			=   processor(text=lowerCamelCase    ,						return_tensors="""np"""		)

      lowerCAmelCase_			:      Tuple			=   tokenizer(lowerCamelCase    ,						return_tensors="""np"""		)

      for key in encoded_tok.keys():
        self.assertListEqual(encoded_tok[key][0].tolist()    ,						encoded_processor[key][0].tolist()		)

    def __lowercase   (   self    :  Tuple		)   ->  List[str]:
      lowerCAmelCase_			:      str			=   self.get_image_processor()
      lowerCAmelCase_			:      Any			=   self.get_tokenizer()

      lowerCAmelCase_			:      Tuple			=   OwlViTProcessor(tokenizer=lowerCamelCase    ,						image_processor=lowerCamelCase		)

      lowerCAmelCase_			:      Union[str, Any]			=   """lower newer"""
      lowerCAmelCase_			:      List[str]			=   self.prepare_image_inputs()

      lowerCAmelCase_			:      Union[str, Any]			=   processor(text=lowerCamelCase    ,						images=lowerCamelCase		)

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

      # test if it raises when no input is passed
      with pytest.raises(lowerCamelCase		):
        processor()

    def __lowercase   (   self    :  Any		)   ->  str:
      lowerCAmelCase_			:      Dict			=   """google/owlvit-base-patch32"""
      lowerCAmelCase_			:      str			=   OwlViTProcessor.from_pretrained(lowerCamelCase		)

      lowerCAmelCase_			:      str			=   ["""cat""", """nasa badge"""]
      lowerCAmelCase_			:      Any			=   processor(text=lowerCamelCase		)

      lowerCAmelCase_			:      List[Any]			=   16
      self.assertListEqual(list(inputs.keys()		)    ,						["""input_ids""", """attention_mask"""]		)
      self.assertEqual(inputs["""input_ids"""].shape    ,						(2, seq_length)		)

      # test if it raises when no input is passed
      with pytest.raises(lowerCamelCase		):
        processor()

    def __lowercase   (   self    :  List[Any]		)   ->  Optional[Any]:
      lowerCAmelCase_			:      Any			=   """google/owlvit-base-patch32"""
      lowerCAmelCase_			:      Optional[int]			=   OwlViTProcessor.from_pretrained(lowerCamelCase		)

      lowerCAmelCase_			:      int			=   [["""cat""", """nasa badge"""], ["""person"""]]
      lowerCAmelCase_			:      List[str]			=   processor(text=lowerCamelCase		)

      lowerCAmelCase_			:      List[Any]			=   16
      lowerCAmelCase_			:      Optional[Any]			=   len(lowerCamelCase		)
      lowerCAmelCase_			:      int			=   max([len(lowerCamelCase		) for texts in input_texts]		)

      self.assertListEqual(list(inputs.keys()		)    ,						["""input_ids""", """attention_mask"""]		)
      self.assertEqual(inputs["""input_ids"""].shape    ,						(batch_size * num_max_text_queries, seq_length)		)

      # test if it raises when no input is passed
      with pytest.raises(lowerCamelCase		):
        processor()

    def __lowercase   (   self    :  int		)   ->  int:
      lowerCAmelCase_			:      Union[str, Any]			=   """google/owlvit-base-patch32"""
      lowerCAmelCase_			:      int			=   OwlViTProcessor.from_pretrained(lowerCamelCase		)

      lowerCAmelCase_			:      List[str]			=   ["""cat""", """nasa badge"""]
      lowerCAmelCase_			:      int			=   processor(text=lowerCamelCase		)

      lowerCAmelCase_			:      int			=   16
      lowerCAmelCase_			:      int			=   inputs["""input_ids"""]
      lowerCAmelCase_			:      List[Any]			=   [
          [4_94_06, 23_68, 4_94_07, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
          [4_94_06, 68_41, 1_13_01, 4_94_07, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
      ]

      self.assertListEqual(list(inputs.keys()		)    ,						["""input_ids""", """attention_mask"""]		)
      self.assertEqual(inputs["""input_ids"""].shape    ,						(2, seq_length)		)
      self.assertListEqual(list(input_ids[0]		)    ,						predicted_ids[0]		)
      self.assertListEqual(list(input_ids[1]		)    ,						predicted_ids[1]		)

    def __lowercase   (   self    :  int		)   ->  Dict:
      lowerCAmelCase_			:      Union[str, Any]			=   self.get_image_processor()
      lowerCAmelCase_			:      Dict			=   self.get_tokenizer()

      lowerCAmelCase_			:      Any			=   OwlViTProcessor(tokenizer=lowerCamelCase    ,						image_processor=lowerCamelCase		)

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

      lowerCAmelCase_			:      List[str]			=   processor(images=lowerCamelCase    ,						query_images=lowerCamelCase		)

      self.assertListEqual(list(inputs.keys()		)    ,						["""query_pixel_values""", """pixel_values"""]		)

      # test if it raises when no input is passed
      with pytest.raises(lowerCamelCase		):
        processor()





    def __lowercase   (   self    :  str		)   ->  Optional[int]:
      lowerCAmelCase_			:      Tuple			=   self.get_image_processor()
      lowerCAmelCase_			:      Dict			=   self.get_tokenizer()

      lowerCAmelCase_			:      Union[str, Any]			=   OwlViTProcessor(tokenizer=lowerCamelCase    ,						image_processor=lowerCamelCase		)

      lowerCAmelCase_			:      int			=   [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

      lowerCAmelCase_			:      List[str]			=   processor.batch_decode(lowerCamelCase		)
      lowerCAmelCase_			:      str			=   tokenizer.batch_decode(lowerCamelCase		)

      self.assertListEqual(lowerCamelCase    ,						lowerCamelCase		)







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







    """simple docstring"""



    def __lowercase   (   self    :  str		)   ->  Any:
      with tempfile.TemporaryDirectory() as tmpdirname:
        # pipeline has Flax weights
        lowerCAmelCase_			:      int			=   FlaxDiffusionPipeline.from_pretrained(
            """hf-internal-testing/tiny-stable-diffusion-pipe"""    ,						safety_checker=lowerCamelCase    ,						cache_dir=lowerCamelCase		)

        lowerCAmelCase_			:      Dict			=   [t[-1] for t in os.walk(os.path.join(lowerCamelCase    ,						os.listdir(lowerCamelCase		)[0]    ,						"""snapshots"""		)		)]
        lowerCAmelCase_			:      List[Any]			=   [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    __snake_case       (   unittest.TestCase):







    """simple docstring"""



    def __lowercase   (   self    :  Any		)   ->  Optional[int]:
      lowerCAmelCase_,							lowerCAmelCase_			:      Union[str, Any]			=   FlaxStableDiffusionPipeline.from_pretrained(
          """hf-internal-testing/tiny-stable-diffusion-pipe"""    ,						safety_checker=lowerCamelCase		)

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

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

      lowerCAmelCase_			:      Dict			=   jax.device_count()
      lowerCAmelCase_			:      Tuple			=   num_samples * [prompt]
      lowerCAmelCase_			:      Optional[int]			=   pipeline.prepare_inputs(lowerCamelCase		)

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

      lowerCAmelCase_			:      str			=   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.1_514_745		) < 1E-3
        assert np.abs(np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 49_947.875		) < 5E-1

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

    def __lowercase   (   self    :  Optional[int]		)   ->  List[str]:
      lowerCAmelCase_,							lowerCAmelCase_			:      Tuple			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""flax"""    ,						safety_checker=lowerCamelCase		)

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

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

      lowerCAmelCase_			:      str			=   jax.device_count()
      lowerCAmelCase_			:      Optional[Any]			=   num_samples * [prompt]
      lowerCAmelCase_			:      Dict			=   pipeline.prepare_inputs(lowerCamelCase		)

      # shard inputs and rng
      lowerCAmelCase_			:      Optional[Any]			=   replicate(lowerCamelCase		)
      lowerCAmelCase_			:      str			=   jax.random.split(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      Union[str, Any]			=   shard(lowerCamelCase		)

      lowerCAmelCase_			:      Optional[int]			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      if jax.device_count() == 8:
        assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    ,						dtype=np.floataa		).sum() - 0.05_652_401)		) < 1E-3
        assert np.abs((np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 2_383_808.2)		) < 5E-1

    def __lowercase   (   self    :  List[Any]		)   ->  List[Any]:
      lowerCAmelCase_,							lowerCAmelCase_			:      Optional[Any]			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""bf16"""    ,						dtype=jnp.bfloataa    ,						safety_checker=lowerCamelCase		)

      lowerCAmelCase_			:      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"""
      )

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

      lowerCAmelCase_			:      Dict			=   jax.device_count()
      lowerCAmelCase_			:      List[Any]			=   num_samples * [prompt]
      lowerCAmelCase_			:      Tuple			=   pipeline.prepare_inputs(lowerCamelCase		)

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

      lowerCAmelCase_			:      Any			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      if jax.device_count() == 8:
        assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    ,						dtype=np.floataa		).sum() - 0.04_003_906)		) < 1E-3
        assert np.abs((np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 2_373_516.75)		) < 5E-1

    def __lowercase   (   self    :  int		)   ->  Optional[int]:
      lowerCAmelCase_,							lowerCAmelCase_			:      int			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""bf16"""    ,						dtype=jnp.bfloataa		)

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

      lowerCAmelCase_			:      str			=   jax.random.PRNGKey(0		)
      lowerCAmelCase_			:      List[Any]			=   50

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

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

      lowerCAmelCase_			:      Tuple			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      if jax.device_count() == 8:
        assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    ,						dtype=np.floataa		).sum() - 0.04_003_906)		) < 1E-3
        assert np.abs((np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 2_373_516.75)		) < 5E-1

    def __lowercase   (   self    :  List[str]		)   ->  Any:
      lowerCAmelCase_			:      List[Any]			=   FlaxDDIMScheduler(
          beta_start=0.00_085    ,						beta_end=0.012    ,						beta_schedule="""scaled_linear"""    ,						set_alpha_to_one=lowerCamelCase    ,						steps_offset=1    ,						)

      lowerCAmelCase_,							lowerCAmelCase_			:      List[str]			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""bf16"""    ,						dtype=jnp.bfloataa    ,						scheduler=lowerCamelCase    ,						safety_checker=lowerCamelCase    ,						)
      lowerCAmelCase_			:      Any			=   scheduler.create_state()

      lowerCAmelCase_			:      Optional[Any]			=   scheduler_state

      lowerCAmelCase_			:      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"""
      )

      lowerCAmelCase_			:      int			=   jax.random.PRNGKey(0		)
      lowerCAmelCase_			:      List[str]			=   50

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

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

      lowerCAmelCase_			:      Tuple			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      if jax.device_count() == 8:
        assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    ,						dtype=np.floataa		).sum() - 0.045_043_945)		) < 1E-3
        assert np.abs((np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 2_347_693.5)		) < 5E-1





    def __lowercase   (   self    :  Any		)   ->  Union[str, Any]:
      lowerCAmelCase_			:      int			=   (
          """A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"""
          """ field, close up, split lighting, cinematic"""
      )

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

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

      lowerCAmelCase_			:      Optional[int]			=   replicate(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   pipeline.prepare_inputs(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   shard(lowerCamelCase		)
      lowerCAmelCase_			:      Tuple			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images
      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      lowerCAmelCase_			:      Dict			=   images[2, 0, 2_56, 10:17, 1]

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

      lowerCAmelCase_			:      List[str]			=   replicate(lowerCamelCase		)
      lowerCAmelCase_			:      Dict			=   pipeline.prepare_inputs(lowerCamelCase		)
      lowerCAmelCase_			:      Union[str, Any]			=   shard(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images
      assert images_eff.shape == (num_samples, 1, 5_12, 5_12, 3)
      lowerCAmelCase_			:      List[Any]			=   images[2, 0, 2_56, 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



