





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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

               __A						:			int								=	[t[-1] for t in os.walk(os.path.join(_UpperCAmelCase						,     os.listdir(_UpperCAmelCase)[0]						,     'snapshots'))]
               __A						:			List[str]								=	[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 			SCREAMING_SNAKE_CASE   (unittest.TestCase							):


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

          __A						:			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'
          )

          __A						:			List[Any]								=	jax.random.PRNGKey(0)
          __A						:			Dict								=	4

          __A						:			Any								=	jax.device_count()
          __A						:			Optional[Any]								=	num_samples * [prompt]
          __A						:			Union[str, Any]								=	pipeline.prepare_inputs(_UpperCAmelCase)

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

          __A						:			Optional[int]								=	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.1514745) < 1e-3
               assert np.abs(np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 49947.875) < 5e-1

          __A						:			List[str]								=	pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))
          assert len(_UpperCAmelCase) == num_samples


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

          __A						:			Optional[Any]								=	(
              'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
              ' field, close up, split lighting, cinematic'
          )

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

          __A						:			Optional[int]								=	jax.device_count()
          __A						:			List[str]								=	num_samples * [prompt]
          __A						:			str								=	pipeline.prepare_inputs(_UpperCAmelCase)

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

          __A						:			List[Any]								=	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.05652401)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 2383808.2)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A			,__A						:			List[str]								=	FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,     revision='bf16'						,     dtype=jnp.bfloataa						,     safety_checker=_UpperCAmelCase)

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

          __A						:			Tuple								=	jax.random.PRNGKey(0)
          __A						:			List[str]								=	50

          __A						:			Dict								=	jax.device_count()
          __A						:			Dict								=	num_samples * [prompt]
          __A						:			str								=	pipeline.prepare_inputs(_UpperCAmelCase)

          # shard inputs and rng
          __A						:			List[str]								=	replicate(_UpperCAmelCase)
          __A						:			Optional[int]								=	jax.random.split(_UpperCAmelCase						,     _UpperCAmelCase)
          __A						:			Any								=	shard(_UpperCAmelCase)

          __A						:			Dict								=	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.04003906)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 2373516.75)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A			,__A						:			Dict								=	FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,     revision='bf16'						,     dtype=jnp.bfloataa)

          __A						:			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'
          )

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

          __A						:			str								=	jax.device_count()
          __A						:			List[str]								=	num_samples * [prompt]
          __A						:			int								=	pipeline.prepare_inputs(_UpperCAmelCase)

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

          __A						:			Any								=	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.04003906)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 2373516.75)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          __A						:			int								=	FlaxDDIMScheduler(
              beta_start=0.00085						,     beta_end=0.012						,     beta_schedule='scaled_linear'						,     set_alpha_to_one=_UpperCAmelCase						,     steps_offset=1						,     )

          __A			,__A						:			Union[str, Any]								=	FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,     revision='bf16'						,     dtype=jnp.bfloataa						,     scheduler=_UpperCAmelCase						,     safety_checker=_UpperCAmelCase						,     )
          __A						:			Tuple								=	scheduler.create_state()

          __A						:			str								=	scheduler_state

          __A						:			Optional[Any]								=	(
              'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
              ' field, close up, split lighting, cinematic'
          )

          __A						:			Optional[Any]								=	jax.random.PRNGKey(0)
          __A						:			List[Any]								=	50

          __A						:			Tuple								=	jax.device_count()
          __A						:			Optional[int]								=	num_samples * [prompt]
          __A						:			int								=	pipeline.prepare_inputs(_UpperCAmelCase)

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

          __A						:			int								=	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.045043945)) < 1e-3
               assert np.abs((np.abs(_UpperCAmelCase						,     dtype=np.floataa).sum() - 2347693.5)) < 5e-1


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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

          __A						:			List[str]								=	jax.device_count()
          __A						:			Optional[Any]								=	num_samples * [prompt]
          __A						:			List[str]								=	jax.random.split(jax.random.PRNGKey(0)						,     _UpperCAmelCase)

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

          __A						:			List[str]								=	replicate(_UpperCAmelCase)
          __A						:			Optional[int]								=	pipeline.prepare_inputs(_UpperCAmelCase)
          __A						:			Optional[Any]								=	shard(_UpperCAmelCase)
          __A						:			int								=	pipeline(_UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     jit=_UpperCAmelCase).images
          assert images.shape == (num_samples, 1, 512, 512, 3)
          __A						:			str								=	images[2, 0, 256, 10:17, 1]

          # With memory efficient attention
          __A			,__A						:			Tuple								=	FlaxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'						,     revision='bf16'						,     dtype=jnp.bfloataa						,     safety_checker=_UpperCAmelCase						,     use_memory_efficient_attention=_UpperCAmelCase						,     )

          __A						:			Optional[int]								=	replicate(_UpperCAmelCase)
          __A						:			int								=	pipeline.prepare_inputs(_UpperCAmelCase)
          __A						:			Tuple								=	shard(_UpperCAmelCase)
          __A						:			Tuple								=	pipeline(_UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     jit=_UpperCAmelCase).images
          assert images_eff.shape == (num_samples, 1, 512, 512, 3)
          __A						:			Any								=	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






'''simple docstring'''




import itertools
import math






def 	_lowerCAmelCase      (  __snake_case  :			int      )					->					bool:

     if 1 < number < 4:
          # 2 and 3 are primes
          return True
     elif number < 2 or number % 2 == 0 or number % 3 == 0:
          # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
          return False

     # All primes number are in format of 6k +/- 1
     for i in range(5			,							int(math.sqrt(__snake_case      ) + 1      )			,							6      ):
          if number % i == 0 or number % (i + 2) == 0:
               return False
     return True






def 	_lowerCAmelCase      (  )					->					List[Any]:
     __A						:			Optional[Any]								=	2
     while True:
          if is_prime(__snake_case      ):
               yield num
          num += 1






def 	_lowerCAmelCase      (  __snake_case  :			int = 1_00_01      )					->					int:
     return next(itertools.islice(prime_generator()			,							nth - 1			,							__snake_case      )      )


if __name__ == "__main__":
  print(f"""{solution() = }""")