


'''simple docstring'''


import gc
import unittest

from diffusers import FlaxStableDiffusionInpaintPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
				import jax
				import jax.numpy as jnp
				from flax.jax_utils import replicate
				from flax.training.common_utils import shard





@slow
@require_flax
class     lowercase__  (						unittest.TestCase					):





							'''simple docstring'''


							def UpperCAmelCase_	(    self  ):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      load_image(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
											    """/sd2-inpaint/init_image.png"""  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      load_image(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"""  )

											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      """xvjiarui/stable-diffusion-2-inpainting"""
											_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE :       Any     		=      FlaxStableDiffusionInpaintPipeline.from_pretrained(__snake_case    ,					safety_checker=__snake_case  )

											_SCREAMING_SNAKE_CASE :       Any     		=      """Face of a yellow cat, high resolution, sitting on a park bench"""

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

											_SCREAMING_SNAKE_CASE :       List[str]     		=      jax.device_count()
											_SCREAMING_SNAKE_CASE :       int     		=      num_samples * [prompt]
											_SCREAMING_SNAKE_CASE :       Any     		=      num_samples * [init_image]
											_SCREAMING_SNAKE_CASE :       str     		=      num_samples * [mask_image]
											_SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE		,       _SCREAMING_SNAKE_CASE :       Tuple     		=      pipeline.prepare_inputs(__snake_case    ,					__snake_case    ,					__snake_case  )

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

											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      pipeline(
											    __snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					jit=__snake_case  )

											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      output.images.reshape(__snake_case    ,					512    ,					512    ,					3  )

											_SCREAMING_SNAKE_CASE :       Tuple     		=      images[0, 253:256, 253:256, -1]

											_SCREAMING_SNAKE_CASE :       Tuple     		=      jnp.asarray(jax.device_get(image_slice.flatten()  )  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      jnp.array(
											    [0.3611307, 0.37649736, 0.3757408, 0.38213953, 0.39295167, 0.3841631, 0.41554978, 0.4137475, 0.4217084]  )
											print(f"""output_slice: {output_slice}"""  )
											assert jnp.abs(output_slice - expected_slice  ).max() < 1e-2








'''simple docstring'''


def 			snake_case_   (  SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       str     		=      0
				for ch in input_str:
								_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      ord(SCREAMING_SNAKE_CASE__     )
								_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      pow(2							,      SCREAMING_SNAKE_CASE__     )

								# If we already turned on bit for current character's unicode
								if bitmap >> ch_unicode & 1 == 1:
												return False
								bitmap |= ch_bit_index_on
				return True


if __name__ == "__main__":
				import doctest

				doctest.testmod()




