


'''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







'''simple docstring'''






def 				UpperCamelCase_			(    A__					:   int					):
  '''simple docstring'''





  lowerCAmelCase_			:      Union[str, Any]			=   (1 + 24 * n) ** 0.5
  return ((1 + root) / 6) % 1 == 0







def 				UpperCamelCase_			(    A__					:   int = 50_00					):
  '''simple docstring'''





  lowerCAmelCase_			:      Optional[int]			=   [(i * (3 * i - 1)) // 2 for i in range(1							,    A__					)]
  for i, pentagonal_i in enumerate(A__					):
    for j in range(A__							,    len(A__					)					):
      lowerCAmelCase_			:      int			=   pentagonal_nums[j]
      lowerCAmelCase_			:      Union[str, Any]			=   pentagonal_i + pentagonal_j
      lowerCAmelCase_			:      List[Any]			=   pentagonal_j - pentagonal_i
      if is_pentagonal(A__					) and is_pentagonal(A__					):
        return b

  return -1


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



