




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


UpperCamelCase__    			=			logging.get_logger(__name__)

UpperCamelCase__    			=			"""▁"""

UpperCamelCase__    			=			{"""vocab_file""": """sentencepiece.bpe.model"""}

UpperCamelCase__    			=			{
    """vocab_file""": {
        """facebook/mbart-large-50-one-to-many-mmt""": (
            """https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model"""
        ),
    }
}

UpperCamelCase__    			=			{
    """facebook/mbart-large-50-one-to-many-mmt""": 1024,
}

# fmt: off
UpperCamelCase__    			=			["""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""", """af_ZA""", """az_AZ""", """bn_IN""", """fa_IR""", """he_IL""", """hr_HR""", """id_ID""", """ka_GE""", """km_KH""", """mk_MK""", """ml_IN""", """mn_MN""", """mr_IN""", """pl_PL""", """ps_AF""", """pt_XX""", """sv_SE""", """sw_KE""", """ta_IN""", """te_IN""", """th_TH""", """tl_XX""", """uk_UA""", """ur_PK""", """xh_ZA""", """gl_ES""", """sl_SI"""]
class 			a__    (					snake_case__      ):
						_a							:    Tuple					=    VOCAB_FILES_NAMES
						_a							:    int					=    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
						_a							:    Optional[int]					=    PRETRAINED_VOCAB_FILES_MAP
						_a							:    Dict					=    ["""input_ids""", """attention_mask"""]

						_a							:    List[int]					=    []
						_a							:    List[int]					=    []






						def __init__(   self		,			_A		,			_A=None		,			_A=None		,			_A="</s>"		,			_A="</s>"		,			_A="<s>"		,			_A="<unk>"		,			_A="<pad>"		,			_A="<mask>"		,			_A = None		,			**_A		,			):




											"""simple docstring"""



											__lowerCAmelCase				   =       AddedToken(_A		,			lstrip=_A		,			rstrip=_A			) if isinstance(_A		,			_A			) else mask_token

											__lowerCAmelCase				   =       {} if sp_model_kwargs is None else sp_model_kwargs

											__lowerCAmelCase				   =       kwargs.get("additional_special_tokens"		,			[]			)
											kwargs["additional_special_tokens"] += [
											    code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
											]

											super().__init__(
											    src_lang=_A		,			tgt_lang=_A		,			eos_token=_A		,			unk_token=_A		,			sep_token=_A		,			cls_token=_A		,			pad_token=_A		,			mask_token=_A		,			sp_model_kwargs=self.sp_model_kwargs		,			**_A		,			)

											__lowerCAmelCase				   =       spm.SentencePieceProcessor(**self.sp_model_kwargs			)
											self.sp_model.Load(str(_A			)			)
											__lowerCAmelCase				   =       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
											__lowerCAmelCase				   =       {"<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
											__lowerCAmelCase				   =       1

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

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

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






						@property
						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""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 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											return self._src_lang






						@src_lang.setter
						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       new_src_lang
											self.set_src_lang_special_tokens(self._src_lang			)






						def __getstate__(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.__dict__.copy()
											__lowerCAmelCase				   =       None
											return state






						def __setstate__(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       d

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

											__lowerCAmelCase				   =       spm.SentencePieceProcessor(**self.sp_model_kwargs			)
											self.sp_model.Load(self.vocab_file			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											if token in self.fairseq_tokens_to_ids:
																return self.fairseq_tokens_to_ids[token]
											__lowerCAmelCase				   =       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 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""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 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       []
											__lowerCAmelCase				   =       ""
											__lowerCAmelCase				   =       False
											for token in tokens:
																# make sure that special tokens are not decoded using sentencepiece model
																if token in self.all_special_tokens:
																					if not prev_is_special:
																										out_string += " "
																					out_string += self.sp_model.decode(_A			) + token
																					__lowerCAmelCase				   =       True
																					__lowerCAmelCase				   =       []
																else:
																					current_sub_tokens.append(_A			)
																					__lowerCAmelCase				   =       False
											out_string += self.sp_model.decode(_A			)
											return out_string.strip()






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A = None			):




											"""simple docstring"""



											if not os.path.isdir(_A			):
																logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""			)
																return
											__lowerCAmelCase				   =       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:
																					__lowerCAmelCase				   =       self.sp_model.serialized_model_proto()
																					fi.write(_A			)

											return (out_vocab_file,)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A = None		,			_A = False			):




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

											__lowerCAmelCase				   =       [1] * len(self.prefix_tokens			)
											__lowerCAmelCase				   =       [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 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A = None			):




											"""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 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A		,			_A		,			**_A			):




											"""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"			)
											__lowerCAmelCase				   =       src_lang
											__lowerCAmelCase				   =       self(_A		,			add_special_tokens=_A		,			return_tensors=_A		,			**_A			)
											__lowerCAmelCase				   =       self.convert_tokens_to_ids(_A			)
											__lowerCAmelCase				   =       tgt_lang_id
											return inputs






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A = "en_XX"		,			_A = None		,			_A = "ro_RO"		,			**_A		,			):




											"""simple docstring"""



											__lowerCAmelCase				   =       src_lang
											__lowerCAmelCase				   =       tgt_lang
											return super().prepare_seqaseq_batch(_A		,			_A		,			**_A			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											return self.set_src_lang_special_tokens(self.src_lang			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											return self.set_tgt_lang_special_tokens(self.tgt_lang			)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.lang_code_to_id[src_lang]
											__lowerCAmelCase				   =       [self.cur_lang_code_id]
											__lowerCAmelCase				   =       [self.eos_token_id]






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.lang_code_to_id[tgt_lang]
											__lowerCAmelCase				   =       [self.cur_lang_code_id]
											__lowerCAmelCase				   =       [self.eos_token_id]












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





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

													_UpperCAmelCase									=			[t[-1] for t in os.walk(os.path.join(UpperCAmelCase       ,    os.listdir(UpperCAmelCase    )[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 __lowerCamelCase (      unittest.TestCase):





	"""simple docstring"""






	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'hf-internal-testing/tiny-stable-diffusion-pipe'       ,    safety_checker=UpperCAmelCase    )

							_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(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

							_UpperCAmelCase									=			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.1_51_47_45    ) < 1e-3
													assert np.abs(np.abs(UpperCAmelCase       ,    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(UpperCAmelCase    ) == num_samples

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='flax'       ,    safety_checker=UpperCAmelCase    )

							_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(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

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

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







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

							_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(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

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

	def        UpperCamelCase    (							self    ):
							"""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(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

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

	def        UpperCamelCase    (							self    ):
							"""simple docstring"""







							_UpperCAmelCase									=			FlaxDDIMScheduler(
							    beta_start=0.0_00_85       ,    beta_end=0.0_12       ,    beta_schedule='scaled_linear'       ,    set_alpha_to_one=UpperCAmelCase       ,    steps_offset=1       ,    )

							_UpperCAmelCase					, _UpperCAmelCase									=			FlaxStableDiffusionPipeline.from_pretrained(
							    'CompVis/stable-diffusion-v1-4'       ,    revision='bf16'       ,    dtype=jnp.bfloataa       ,    scheduler=UpperCAmelCase       ,    safety_checker=UpperCAmelCase       ,    )
							_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(UpperCAmelCase    )

							# shard inputs and rng
							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			jax.random.split(UpperCAmelCase       ,    UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )

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

	def        UpperCamelCase    (							self    ):
							"""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    )       ,    UpperCAmelCase    )

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

							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).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=UpperCAmelCase       ,    use_memory_efficient_attention=UpperCAmelCase       ,    )

							_UpperCAmelCase									=			replicate(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline.prepare_inputs(UpperCAmelCase    )
							_UpperCAmelCase									=			shard(UpperCAmelCase    )
							_UpperCAmelCase									=			pipeline(UpperCAmelCase       ,    UpperCAmelCase       ,    UpperCAmelCase       ,    jit=UpperCAmelCase    ).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
