





import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


__lowercase  :   Tuple						      =logging.get_logger(__name__)

__lowercase  :   Tuple						      ="""▁"""

__lowercase  :   Optional[Any]						      ={"""vocab_file""": """sentencepiece.bpe.model"""}

__lowercase  :   List[str]						      ={
    """vocab_file""": {
        """facebook/mbart-large-en-ro""": (
            """https://huggingface.co/facebook/mbart-large-en-ro/resolve/main/sentencepiece.bpe.model"""
        ),
        """facebook/mbart-large-cc25""": (
            """https://huggingface.co/facebook/mbart-large-cc25/resolve/main/sentencepiece.bpe.model"""
        ),
    }
}

__lowercase  :   Dict						      ={
    """facebook/mbart-large-en-ro""": 1024,
    """facebook/mbart-large-cc25""": 1024,
}

# fmt: off
__lowercase  :   int						      =["""ar_AR""", """cs_CZ""", """de_DE""", """en_XX""", """es_XX""", """et_EE""", """fi_FI""", """fr_XX""", """gu_IN""", """hi_IN""", """it_IT""", """ja_XX""", """kk_KZ""", """ko_KR""", """lt_LT""", """lv_LV""", """my_MM""", """ne_NP""", """nl_XX""", """ro_RO""", """ru_RU""", """si_LK""", """tr_TR""", """vi_VN""", """zh_CN"""]

class      A       (	__snake_case						):
      _snake_case							     =VOCAB_FILES_NAMES
      _snake_case							     =PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _snake_case							     =PRETRAINED_VOCAB_FILES_MAP
      _snake_case							     =['input_ids', 'attention_mask']

      _snake_case							     =[]
      _snake_case							     =[]



      def __init__(	self:						int		,  _lowerCAmelCase:						List[str]		,  _lowerCAmelCase:						int="<s>"		,  _lowerCAmelCase:						Optional[Any]="</s>"		,  _lowerCAmelCase:						Optional[int]="</s>"		,  _lowerCAmelCase:						Optional[Any]="<s>"		,  _lowerCAmelCase:						str="<unk>"		,  _lowerCAmelCase:						str="<pad>"		,  _lowerCAmelCase:						Union[str, Any]="<mask>"		,  _lowerCAmelCase:						int=None		,  _lowerCAmelCase:						List[Any]=None		,  _lowerCAmelCase:						int=None		,  _lowerCAmelCase:						Optional[int] = None		,  _lowerCAmelCase:						str=None		,  **_lowerCAmelCase:						Optional[Any]		,  )						->    List[Any]:
         '''simple docstring'''

         UpperCAmelCase_											=AddedToken(A_		,  lstrip=A_		,  rstrip=A_       ) if isinstance(A_		,  A_       ) else mask_token

         UpperCAmelCase_											={} if sp_model_kwargs is None else sp_model_kwargs

         super().__init__(
             bos_token=A_		,  eos_token=A_		,  unk_token=A_		,  sep_token=A_		,  cls_token=A_		,  pad_token=A_		,  mask_token=A_		,  tokenizer_file=A_		,  src_lang=A_		,  tgt_lang=A_		,  additional_special_tokens=A_		,  sp_model_kwargs=self.sp_model_kwargs		,  **A_		,  )

         UpperCAmelCase_											=spm.SentencePieceProcessor(**self.sp_model_kwargs       )
         self.sp_model.Load(str(A_       )       )
         UpperCAmelCase_											=vocab_file

         # Original fairseq vocab and spm vocab must be "aligned":
         # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
         # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
         # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
         # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

         # Mimic fairseq token-to-id alignment for the first 4 token
         UpperCAmelCase_											={"<s>": 0, "<pad>": 1, "</s>": 2, "<unk>": 3}

         # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
         UpperCAmelCase_											=1

         UpperCAmelCase_											=len(self.sp_model       )
         UpperCAmelCase_											={
             code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(A_       )
         }
         UpperCAmelCase_											={v: k for k, v in self.lang_code_to_id.items()}
         UpperCAmelCase_											=len(self.sp_model       ) + len(self.lang_code_to_id       ) + self.fairseq_offset

         self.fairseq_tokens_to_ids.update(self.lang_code_to_id       )
         UpperCAmelCase_											={v: k for k, v in self.fairseq_tokens_to_ids.items()}
         UpperCAmelCase_											=list(self.lang_code_to_id.keys()       )

         if additional_special_tokens is not None:
            # Only add those special tokens if they are not already there.
            self._additional_special_tokens.extend(
                [t for t in additional_special_tokens if t not in self._additional_special_tokens]       )

         UpperCAmelCase_											=src_lang if src_lang is not None else "en_XX"
         UpperCAmelCase_											=self.lang_code_to_id[self._src_lang]
         UpperCAmelCase_											=tgt_lang
         self.set_src_lang_special_tokens(self._src_lang       )



      def __getstate__(	self:						Optional[Any]       )						->    List[Any]:
         '''simple docstring'''

         UpperCAmelCase_											=self.__dict__.copy()
         UpperCAmelCase_											=None
         UpperCAmelCase_											=self.sp_model.serialized_model_proto()
         return state



      def __setstate__(	self:						List[Any]		,  _lowerCAmelCase:						str       )						->    Union[str, Any]:
         '''simple docstring'''

         UpperCAmelCase_											=d

         # for backward compatibility
         if not hasattr(self		,  "sp_model_kwargs"       ):
            UpperCAmelCase_											={}

         UpperCAmelCase_											=spm.SentencePieceProcessor(**self.sp_model_kwargs       )
         self.sp_model.LoadFromSerializedProto(self.sp_model_proto       )



      @property
      def 	lowerCAmelCase__    (	self:						List[str]       )						->    Optional[int]:
         '''simple docstring'''

         return len(self.sp_model       ) + len(self.lang_code_to_id       ) + self.fairseq_offset + 1  # Plus 1 for the mask token



      @property
      def 	lowerCAmelCase__    (	self:						str       )						->    Tuple:
         '''simple docstring'''

         return self._src_lang



      @src_lang.setter
      def 	lowerCAmelCase__    (	self:						Optional[int]		,  _lowerCAmelCase:						Any       )						->    Any:
         '''simple docstring'''

         UpperCAmelCase_											=new_src_lang
         self.set_src_lang_special_tokens(self._src_lang       )



      def 	lowerCAmelCase__    (	self:						Optional[int]		,  _lowerCAmelCase:						Union[str, Any]		,  _lowerCAmelCase:						List[Any] = None		,  _lowerCAmelCase:						Union[str, Any] = False       )						->    int:
         '''simple docstring'''

         if already_has_special_tokens:
            return super().get_special_tokens_mask(
                token_ids_a=A_		,  token_ids_a=A_		,  already_has_special_tokens=A_       )

         UpperCAmelCase_											=[1] * len(self.prefix_tokens       )
         UpperCAmelCase_											=[1] * len(self.suffix_tokens       )
         if token_ids_a is None:
            return prefix_ones + ([0] * len(A_       )) + suffix_ones
         return prefix_ones + ([0] * len(A_       )) + ([0] * len(A_       )) + suffix_ones



      def 	lowerCAmelCase__    (	self:						Optional[Any]		,  _lowerCAmelCase:						Dict		,  _lowerCAmelCase:						Dict = None       )						->    Any:
         '''simple docstring'''

         if token_ids_a is None:
            return self.prefix_tokens + token_ids_a + self.suffix_tokens
         # We don't expect to process pairs, but leave the pair logic for API consistency
         return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens



      def 	lowerCAmelCase__    (	self:						List[Any]		,  _lowerCAmelCase:						Optional[int]		,  _lowerCAmelCase:						str = None       )						->    Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_											=[self.sep_token_id]
         UpperCAmelCase_											=[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 + sep + token_ids_a + sep       ) * [0]



      def 	lowerCAmelCase__    (	self:						Any		,  _lowerCAmelCase:						Tuple		,  _lowerCAmelCase:						List[str]		,  _lowerCAmelCase:						int		,  _lowerCAmelCase:						int		,  **_lowerCAmelCase:						Union[str, Any]       )						->    List[Any]:
         '''simple docstring'''

         if src_lang is None or tgt_lang is None:
            raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model"       )
         UpperCAmelCase_											=src_lang
         UpperCAmelCase_											=self(A_		,  add_special_tokens=A_		,  return_tensors=A_		,  **A_       )
         UpperCAmelCase_											=self.convert_tokens_to_ids(A_       )
         UpperCAmelCase_											=tgt_lang_id
         return inputs



      def 	lowerCAmelCase__    (	self:						List[str]       )						->    Tuple:
         '''simple docstring'''

         UpperCAmelCase_											={self.convert_ids_to_tokens(A_       ): i for i in range(self.vocab_size       )}
         vocab.update(self.added_tokens_encoder       )
         return vocab



      def 	lowerCAmelCase__    (	self:						str		,  _lowerCAmelCase:						Union[str, Any]       )						->    Tuple:
         '''simple docstring'''

         return self.sp_model.encode(A_		,  out_type=A_       )



      def 	lowerCAmelCase__    (	self:						Tuple		,  _lowerCAmelCase:						Dict       )						->    Any:
         '''simple docstring'''

         if token in self.fairseq_tokens_to_ids:
            return self.fairseq_tokens_to_ids[token]
         UpperCAmelCase_											=self.sp_model.PieceToId(A_       )

         # Need to return unknown token if the SP model returned 0
         return spm_id + self.fairseq_offset if spm_id else self.unk_token_id



      def 	lowerCAmelCase__    (	self:						str		,  _lowerCAmelCase:						Optional[int]       )						->    List[str]:
         '''simple docstring'''

         if index in self.fairseq_ids_to_tokens:
            return self.fairseq_ids_to_tokens[index]
         return self.sp_model.IdToPiece(index - self.fairseq_offset       )



      def 	lowerCAmelCase__    (	self:						List[str]		,  _lowerCAmelCase:						Any       )						->    Union[str, Any]:
         '''simple docstring'''

         UpperCAmelCase_											="".join(A_       ).replace(A_		,  " "       ).strip()
         return out_string



      def 	lowerCAmelCase__    (	self:						List[str]		,  _lowerCAmelCase:						Union[str, Any]		,  _lowerCAmelCase:						Dict = None       )						->    int:
         '''simple docstring'''

         if not os.path.isdir(A_       ):
            logger.error(F'Vocabulary path ({save_directory}) should be a directory'       )
            return
         UpperCAmelCase_											=os.path.join(
             A_		,  (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]       )

         if os.path.abspath(self.vocab_file       ) != os.path.abspath(A_       ) and os.path.isfile(self.vocab_file       ):
            copyfile(self.vocab_file		,  A_       )
         elif not os.path.isfile(self.vocab_file       ):
            with open(A_		,  "wb"       ) as fi:
               UpperCAmelCase_											=self.sp_model.serialized_model_proto()
               fi.write(A_       )

         return (out_vocab_file,)



      def 	lowerCAmelCase__    (	self:						Dict		,  _lowerCAmelCase:						Optional[int]		,  _lowerCAmelCase:						List[str] = "en_XX"		,  _lowerCAmelCase:						int = None		,  _lowerCAmelCase:						str = "ro_RO"		,  **_lowerCAmelCase:						Any		,  )						->    Tuple:
         '''simple docstring'''

         UpperCAmelCase_											=src_lang
         UpperCAmelCase_											=tgt_lang
         return super().prepare_seqaseq_batch(A_		,  A_		,  **A_       )



      def 	lowerCAmelCase__    (	self:						Union[str, Any]       )						->    int:
         '''simple docstring'''

         return self.set_src_lang_special_tokens(self.src_lang       )



      def 	lowerCAmelCase__    (	self:						Union[str, Any]       )						->    Optional[Any]:
         '''simple docstring'''

         return self.set_tgt_lang_special_tokens(self.tgt_lang       )



      def 	lowerCAmelCase__    (	self:						Dict		,  _lowerCAmelCase:						str       )						->    Union[str, Any]:
         '''simple docstring'''

         UpperCAmelCase_											=self.lang_code_to_id[src_lang]
         UpperCAmelCase_											=[]
         UpperCAmelCase_											=[self.eos_token_id, self.cur_lang_code]






      def 	lowerCAmelCase__    (	self:						Optional[Any]		,  _lowerCAmelCase:						List[Any]       )						->    List[str]:
         '''simple docstring'''

         UpperCAmelCase_											=self.lang_code_to_id[lang]
         UpperCAmelCase_											=[]
         UpperCAmelCase_											=[self.eos_token_id, self.cur_lang_code]







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



      def 	lowerCAmelCase__    (	self:						str       )						->    Union[str, Any]:
         '''simple docstring'''

         with tempfile.TemporaryDirectory() as tmpdirname:
            # pipeline has Flax weights
            UpperCAmelCase_											=FlaxDiffusionPipeline.from_pretrained(
                "hf-internal-testing/tiny-stable-diffusion-pipe"		,  safety_checker=_lowerCAmelCase		,  cache_dir=_lowerCAmelCase       )

            UpperCAmelCase_											=[t[-1] for t in os.walk(os.path.join(_lowerCAmelCase		,  os.listdir(_lowerCAmelCase       )[0]		,  "snapshots"       )       )]
            UpperCAmelCase_											=[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      A       (	unittest.TestCase						):



      def 	lowerCAmelCase__    (	self:						Union[str, Any]       )						->    Tuple:
         '''simple docstring'''

         UpperCAmelCase_						,       UpperCAmelCase_											=FlaxStableDiffusionPipeline.from_pretrained(
             "hf-internal-testing/tiny-stable-diffusion-pipe"		,  safety_checker=_lowerCAmelCase       )

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=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_51_47_45       ) < 1e-3
            assert np.abs(np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 4_99_47.8_75       ) < 5e-1

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



      def 	lowerCAmelCase__    (	self:						List[Any]       )						->    Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_						,       UpperCAmelCase_											=FlaxStableDiffusionPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4"		,  revision="flax"		,  safety_checker=_lowerCAmelCase       )

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=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.05_65_24_01)       ) < 1e-3
            assert np.abs((np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 2_38_38_08.2)       ) < 5e-1



      def 	lowerCAmelCase__    (	self:						int       )						->    Union[str, Any]:
         '''simple docstring'''

         UpperCAmelCase_						,       UpperCAmelCase_											=FlaxStableDiffusionPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4"		,  revision="bf16"		,  dtype=jnp.bfloataa		,  safety_checker=_lowerCAmelCase       )

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=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.04_00_39_06)       ) < 1e-3
            assert np.abs((np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 2_37_35_16.75)       ) < 5e-1



      def 	lowerCAmelCase__    (	self:						Optional[int]       )						->    Tuple:
         '''simple docstring'''

         UpperCAmelCase_						,       UpperCAmelCase_											=FlaxStableDiffusionPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4"		,  revision="bf16"		,  dtype=jnp.bfloataa       )

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=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.04_00_39_06)       ) < 1e-3
            assert np.abs((np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 2_37_35_16.75)       ) < 5e-1



      def 	lowerCAmelCase__    (	self:						Optional[Any]       )						->    str:
         '''simple docstring'''

         UpperCAmelCase_											=FlaxDDIMScheduler(
             beta_start=0.0_00_85		,  beta_end=0.0_12		,  beta_schedule="scaled_linear"		,  set_alpha_to_one=_lowerCAmelCase		,  steps_offset=1		,  )

         UpperCAmelCase_						,       UpperCAmelCase_											=FlaxStableDiffusionPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4"		,  revision="bf16"		,  dtype=jnp.bfloataa		,  scheduler=_lowerCAmelCase		,  safety_checker=_lowerCAmelCase		,  )
         UpperCAmelCase_											=scheduler.create_state()

         UpperCAmelCase_											=scheduler_state

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

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

         UpperCAmelCase_											=jax.device_count()
         UpperCAmelCase_											=num_samples * [prompt]
         UpperCAmelCase_											=pipeline.prepare_inputs(_lowerCAmelCase       )

         # shard inputs and rng
         UpperCAmelCase_											=replicate(_lowerCAmelCase       )
         UpperCAmelCase_											=jax.random.split(_lowerCAmelCase		,  _lowerCAmelCase       )
         UpperCAmelCase_											=shard(_lowerCAmelCase       )

         UpperCAmelCase_											=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.0_45_04_39_45)       ) < 1e-3
            assert np.abs((np.abs(_lowerCAmelCase		,  dtype=np.floataa       ).sum() - 2_34_76_93.5)       ) < 5e-1






      def 	lowerCAmelCase__    (	self:						Optional[Any]       )						->    int:
         '''simple docstring'''

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

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

         UpperCAmelCase_						,       UpperCAmelCase_											=FlaxStableDiffusionPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4"		,  revision="bf16"		,  dtype=jnp.bfloataa		,  safety_checker=_lowerCAmelCase		,  )

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

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

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




