


'''simple docstring'''






import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import BatchEncoding, MarianTokenizer
from transformers.testing_utils import get_tests_dir, require_sentencepiece, slow
from transformers.utils import is_sentencepiece_available, is_tf_available, is_torch_available


if is_sentencepiece_available():
  from transformers.models.marian.tokenization_marian import VOCAB_FILES_NAMES, save_json

from ...test_tokenization_common import TokenizerTesterMixin


__A					:  Tuple											=      get_tests_dir("fixtures/test_sentencepiece.model")

__A					:  Union[str, Any]											=      {"target_lang": "fi", "source_lang": "en"}
__A					:  Any											=      ">>zh<<"
__A					:  Any											=      "Helsinki-NLP/"

if is_torch_available():
  __A					:  Union[str, Any]											=      "pt"
elif is_tf_available():
  __A					:  Tuple											=      "tf"
else:
  __A					:  Optional[int]											=      "jax"






@require_sentencepiece
class    __snake_case       (   _SCREAMING_SNAKE_CASE	,unittest.TestCase):







    """simple docstring"""



    lowercase    =					MarianTokenizer
    lowercase    =					False
    lowercase    =					True

    def __lowercase   (   self    :  Any		)   ->  Optional[Any]:
      super().setUp()
      lowerCAmelCase_			:      Dict			=   ["""</s>""", """<unk>""", """▁This""", """▁is""", """▁a""", """▁t""", """est""", """\u0120""", """<pad>"""]
      lowerCAmelCase_			:      List[Any]			=   dict(zip(lowerCamelCase    ,						range(len(lowerCamelCase		)		)		)		)
      lowerCAmelCase_			:      Union[str, Any]			=   Path(self.tmpdirname		)
      save_json(lowerCamelCase    ,						save_dir / VOCAB_FILES_NAMES["""vocab"""]		)
      save_json(lowerCamelCase    ,						save_dir / VOCAB_FILES_NAMES["""tokenizer_config_file"""]		)
      if not (save_dir / VOCAB_FILES_NAMES["source_spm"]).exists():
        copyfile(lowerCamelCase    ,						save_dir / VOCAB_FILES_NAMES["""source_spm"""]		)
        copyfile(lowerCamelCase    ,						save_dir / VOCAB_FILES_NAMES["""target_spm"""]		)

      lowerCAmelCase_			:      int			=   MarianTokenizer.from_pretrained(self.tmpdirname		)
      tokenizer.save_pretrained(self.tmpdirname		)

    def __lowercase   (   self    :  Dict    ,						**lowerCamelCase    :  Optional[Any]		)   ->  MarianTokenizer:
      return MarianTokenizer.from_pretrained(self.tmpdirname    ,						**lowerCamelCase		)

    def __lowercase   (   self    :  Optional[Any]    ,						lowerCamelCase    :  List[Any]		)   ->  Optional[int]:
      return (
          "This is a test",
          "This is a test",
      )

    def __lowercase   (   self    :  Optional[Any]		)   ->  str:
      lowerCAmelCase_			:      Optional[int]			=   """</s>"""
      lowerCAmelCase_			:      Union[str, Any]			=   0

      self.assertEqual(self.get_tokenizer()._convert_token_to_id(lowerCamelCase		)    ,						lowerCamelCase		)
      self.assertEqual(self.get_tokenizer()._convert_id_to_token(lowerCamelCase		)    ,						lowerCamelCase		)

    def __lowercase   (   self    :  Optional[int]		)   ->  Optional[int]:
      lowerCAmelCase_			:      str			=   list(self.get_tokenizer().get_vocab().keys()		)

      self.assertEqual(vocab_keys[0]    ,						"""</s>"""		)
      self.assertEqual(vocab_keys[1]    ,						"""<unk>"""		)
      self.assertEqual(vocab_keys[-1]    ,						"""<pad>"""		)
      self.assertEqual(len(lowerCamelCase		)    ,						9		)

    def __lowercase   (   self    :  int		)   ->  Tuple:
      self.assertEqual(self.get_tokenizer().vocab_size    ,						9		)

    def __lowercase   (   self    :  Union[str, Any]		)   ->  Tuple:
      lowerCAmelCase_			:      Union[str, Any]			=   MarianTokenizer.from_pretrained(F'{ORG_NAME}opus-mt-en-de'		)
      lowerCAmelCase_			:      List[str]			=   en_de_tokenizer(["""I am a small frog"""]    ,						return_tensors=lowerCamelCase		)
      self.assertIsInstance(lowerCamelCase    ,						lowerCamelCase		)
      lowerCAmelCase_			:      Any			=   [38, 1_21, 14, 6_97, 3_88_48, 0]
      self.assertListEqual(lowerCamelCase    ,						batch.input_ids[0]		)

      lowerCAmelCase_			:      Union[str, Any]			=   tempfile.mkdtemp()
      en_de_tokenizer.save_pretrained(lowerCamelCase		)
      lowerCAmelCase_			:      Dict			=   [x.name for x in Path(lowerCamelCase		).glob("""*"""		)]
      self.assertIn("""source.spm"""    ,						lowerCamelCase		)
      MarianTokenizer.from_pretrained(lowerCamelCase		)

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

      lowerCAmelCase_			:      List[str]			=   tok(
          ["""I am a small frog""" * 10_00, """I am a small frog"""]    ,						padding=lowerCamelCase    ,						truncation=lowerCamelCase    ,						return_tensors=lowerCamelCase		)
      self.assertIsInstance(lowerCamelCase    ,						lowerCamelCase		)
      self.assertEqual(batch.input_ids.shape    ,						(2, 5_12)		)

    def __lowercase   (   self    :  Optional[Any]		)   ->  Union[str, Any]:
      lowerCAmelCase_			:      Dict			=   self.get_tokenizer()
      lowerCAmelCase_			:      str			=   tok(["""I am a tiny frog""", """I am a small frog"""]    ,						padding=lowerCamelCase    ,						return_tensors=lowerCamelCase		)
      self.assertIsInstance(lowerCamelCase    ,						lowerCamelCase		)
      self.assertEqual(batch_smaller.input_ids.shape    ,						(2, 10)		)

    @slow
    def __lowercase   (   self    :  Any		)   ->  str:
      # fmt: off
      lowerCAmelCase_			:      List[Any]			=   {"""input_ids""": [[4_34_95, 4_62, 20, 4_21_64, 13_69, 52, 4_64, 1_32, 17_03, 4_92, 13, 74_91, 3_89_99, 6, 8, 4_64, 1_32, 17_03, 4_92, 13, 46_69, 3_78_67, 13, 75_25, 27, 15_93, 9_88, 13, 3_39_72, 70_29, 6, 20, 82_51, 3_83, 2, 2_70, 58_66, 37_88, 2, 23_53, 82_51, 1_23_38, 2, 1_39_58, 3_87, 2, 36_29, 69_53, 1_88, 29_00, 2, 1_39_58, 80_11, 1_15_01, 23, 84_60, 40_73, 3_40_09, 20, 4_35, 1_14_39, 27, 8, 84_60, 40_73, 60_04, 20, 99_88, 3_75, 27, 33, 2_66, 19_45, 10_76, 13_50, 3_78_67, 32_88, 5, 5_77, 10_76, 43_74, 8, 50_82, 5, 2_64_53, 2_57, 5_56, 4_03, 2, 2_42, 1_32, 3_83, 3_16, 4_92, 8, 1_07_67, 6, 3_16, 3_04, 42_39, 3, 0], [1_48, 1_57_22, 19, 18_39, 12, 13_50, 13, 2_23_27, 50_82, 54_18, 4_75_67, 3_59_38, 59, 3_18, 1_95_52, 1_08, 21_83, 54, 1_49_76, 48_35, 32, 5_47, 11_14, 8, 3_15, 24_17, 5, 92, 1_90_88, 3, 0, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00], [36, 63_95, 1_25_70, 3_91_47, 1_15_97, 6, 2_66, 4, 4_54_05, 72_96, 3, 0, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00]], """attention_mask""": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
      # fmt: on

      self.tokenizer_integration_test_util(
          expected_encoding=lowerCamelCase    ,						model_name="""Helsinki-NLP/opus-mt-en-de"""    ,						revision="""1a8c2263da11e68e50938f97e10cd57820bd504c"""    ,						decode_kwargs={"""use_source_tokenizer""": True}    ,						)





    def __lowercase   (   self    :  List[str]		)   ->  str:
      lowerCAmelCase_			:      List[str]			=   MarianTokenizer.from_pretrained("""hf-internal-testing/test-marian-two-vocabs"""		)

      lowerCAmelCase_			:      str			=   """Tämä on testi"""
      lowerCAmelCase_			:      int			=   """This is a test"""

      lowerCAmelCase_			:      Tuple			=   [76, 7, 20_47, 2]
      lowerCAmelCase_			:      Any			=   [69, 12, 11, 9_40, 2]

      lowerCAmelCase_			:      List[str]			=   tokenizer(lowerCamelCase		).input_ids
      self.assertListEqual(lowerCamelCase    ,						lowerCamelCase		)

      lowerCAmelCase_			:      Any			=   tokenizer(text_target=lowerCamelCase		).input_ids
      self.assertListEqual(lowerCamelCase    ,						lowerCamelCase		)

      lowerCAmelCase_			:      str			=   tokenizer.decode(lowerCamelCase    ,						skip_special_tokens=lowerCamelCase		)
      self.assertEqual(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



