


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







   def 			lowerCamelCase		(     self						)     ->		Tuple:

       '''simple docstring'''




       with tempfile.TemporaryDirectory() as tmpdirname:
           # pipeline has Flax weights
           snake_case   :				Union[str, Any]				 =   FlaxDiffusionPipeline.from_pretrained(
               "hf-internal-testing/tiny-stable-diffusion-pipe"						,     safety_checker=_lowercase						,     cache_dir=_lowercase						)

           snake_case   :				str				 =   [t[-1] for t in os.walk(os.path.join(_lowercase						,     os.listdir(_lowercase						)[0]						,     "snapshots"						)						)]
           snake_case   :				Optional[int]				 =   [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						):







   def 			lowerCamelCase		(     self						)     ->		Dict:

       '''simple docstring'''




       snake_case ,snake_case   :				Any				 =   FlaxStableDiffusionPipeline.from_pretrained(
           "hf-internal-testing/tiny-stable-diffusion-pipe"						,     safety_checker=_lowercase						)

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

       snake_case   :				str				 =   jax.random.PRNGKey(0						)
       snake_case   :				Union[str, Any]				 =   4

       snake_case   :				str				 =   jax.device_count()
       snake_case   :				Any				 =   num_samples * [prompt]
       snake_case   :				Any				 =   pipeline.prepare_inputs(_lowercase						)

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

       snake_case   :				Any				 =   pipeline(_lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     jit=_lowercase						).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(_lowercase						,     dtype=np.floataa						).sum() - 4_9947.875						) < 5e-1

       snake_case   :				Optional[Any]				 =   pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]						)						)						)
       assert len(_lowercase						) == num_samples







   def 			lowerCamelCase		(     self						)     ->		Union[str, Any]:

       '''simple docstring'''




       snake_case ,snake_case   :				Optional[int]				 =   FlaxStableDiffusionPipeline.from_pretrained(
           "CompVis/stable-diffusion-v1-4"						,     revision="flax"						,     safety_checker=_lowercase						)

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

       snake_case   :				Optional[Any]				 =   jax.random.PRNGKey(0						)
       snake_case   :				Optional[int]				 =   50

       snake_case   :				int				 =   jax.device_count()
       snake_case   :				Any				 =   num_samples * [prompt]
       snake_case   :				Tuple				 =   pipeline.prepare_inputs(_lowercase						)

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

       snake_case   :				str				 =   pipeline(_lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     jit=_lowercase						).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(_lowercase						,     dtype=np.floataa						).sum() - 238_3808.2)						) < 5e-1







   def 			lowerCamelCase		(     self						)     ->		List[str]:

       '''simple docstring'''




       snake_case ,snake_case   :				Dict				 =   FlaxStableDiffusionPipeline.from_pretrained(
           "CompVis/stable-diffusion-v1-4"						,     revision="bf16"						,     dtype=jnp.bfloataa						,     safety_checker=_lowercase						)

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

       snake_case   :				Union[str, Any]				 =   jax.random.PRNGKey(0						)
       snake_case   :				List[str]				 =   50

       snake_case   :				int				 =   jax.device_count()
       snake_case   :				Optional[Any]				 =   num_samples * [prompt]
       snake_case   :				Tuple				 =   pipeline.prepare_inputs(_lowercase						)

       # shard inputs and rng
       snake_case   :				Tuple				 =   replicate(_lowercase						)
       snake_case   :				Any				 =   jax.random.split(_lowercase						,     _lowercase						)
       snake_case   :				Tuple				 =   shard(_lowercase						)

       snake_case   :				List[Any]				 =   pipeline(_lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     jit=_lowercase						).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(_lowercase						,     dtype=np.floataa						).sum() - 237_3516.75)						) < 5e-1







   def 			lowerCamelCase		(     self						)     ->		List[Any]:

       '''simple docstring'''




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

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

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

       snake_case   :				int				 =   jax.device_count()
       snake_case   :				Optional[int]				 =   num_samples * [prompt]
       snake_case   :				Optional[Any]				 =   pipeline.prepare_inputs(_lowercase						)

       # shard inputs and rng
       snake_case   :				Tuple				 =   replicate(_lowercase						)
       snake_case   :				Any				 =   jax.random.split(_lowercase						,     _lowercase						)
       snake_case   :				Union[str, Any]				 =   shard(_lowercase						)

       snake_case   :				Optional[Any]				 =   pipeline(_lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     jit=_lowercase						).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(_lowercase						,     dtype=np.floataa						).sum() - 237_3516.75)						) < 5e-1







   def 			lowerCamelCase		(     self						)     ->		Any:

       '''simple docstring'''




       snake_case   :				Optional[Any]				 =   FlaxDDIMScheduler(
           beta_start=0.00085						,     beta_end=0.012						,     beta_schedule="scaled_linear"						,     set_alpha_to_one=_lowercase						,     steps_offset=1						,     )

       snake_case ,snake_case   :				Union[str, Any]				 =   FlaxStableDiffusionPipeline.from_pretrained(
           "CompVis/stable-diffusion-v1-4"						,     revision="bf16"						,     dtype=jnp.bfloataa						,     scheduler=_lowercase						,     safety_checker=_lowercase						,     )
       snake_case   :				int				 =   scheduler.create_state()

       snake_case   :				Any				 =   scheduler_state

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

       snake_case   :				Tuple				 =   jax.random.PRNGKey(0						)
       snake_case   :				List[Any]				 =   50

       snake_case   :				Dict				 =   jax.device_count()
       snake_case   :				List[str]				 =   num_samples * [prompt]
       snake_case   :				Union[str, Any]				 =   pipeline.prepare_inputs(_lowercase						)

       # shard inputs and rng
       snake_case   :				Optional[int]				 =   replicate(_lowercase						)
       snake_case   :				str				 =   jax.random.split(_lowercase						,     _lowercase						)
       snake_case   :				Optional[int]				 =   shard(_lowercase						)

       snake_case   :				Dict				 =   pipeline(_lowercase						,     _lowercase						,     _lowercase						,     _lowercase						,     jit=_lowercase						).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(_lowercase						,     dtype=np.floataa						).sum() - 234_7693.5)						) < 5e-1







   def 			lowerCamelCase		(     self						)     ->		Dict:

       '''simple docstring'''




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

       snake_case   :				Union[str, Any]				 =   jax.device_count()
       snake_case   :				Any				 =   num_samples * [prompt]
       snake_case   :				str				 =   jax.random.split(jax.random.PRNGKey(0						)						,     _lowercase						)

       snake_case ,snake_case   :				int				 =   FlaxStableDiffusionPipeline.from_pretrained(
           "CompVis/stable-diffusion-v1-4"						,     revision="bf16"						,     dtype=jnp.bfloataa						,     safety_checker=_lowercase						,     )

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

       # With memory efficient attention
       snake_case ,snake_case   :				Optional[int]				 =   FlaxStableDiffusionPipeline.from_pretrained(
           "CompVis/stable-diffusion-v1-4"						,     revision="bf16"						,     dtype=jnp.bfloataa						,     safety_checker=_lowercase						,     use_memory_efficient_attention=_lowercase						,     )

       snake_case   :				List[Any]				 =   replicate(_lowercase						)
       snake_case   :				str				 =   pipeline.prepare_inputs(_lowercase						)
       snake_case   :				List[str]				 =   shard(_lowercase						)
       snake_case   :				Optional[Any]				 =   pipeline(_lowercase						,     _lowercase						,     _lowercase						,     jit=_lowercase						).images
       assert images_eff.shape == (num_samples, 1, 512, 512, 3)
       snake_case   :				Optional[int]				 =   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 math
import sys







def   __lowerCAmelCase						(						lowercase  :  int			)     ->						int:


    """simple docstring"""





    if number != int(lowercase			):
        raise ValueError("the value of input must be a natural number"			)
    if number < 0:
        raise ValueError("the value of input must not be a negative number"			)
    if number == 0:
        return 1
    snake_case   :				Optional[Any]				 =   [-1] * (number + 1)
    snake_case   :				str				 =   0
    for i in range(1						,							number + 1			):
        snake_case   :				List[Any]				 =   sys.maxsize
        snake_case   :				Union[str, Any]				 =   int(math.sqrt(lowercase			)			)
        for j in range(1						,							root + 1			):
            snake_case   :				List[str]				 =   1 + answers[i - (j**2)]
            snake_case   :				Optional[Any]				 =   min(lowercase						,							lowercase			)
        snake_case   :				Any				 =   answer
    return answers[number]


if __name__ == "__main__":
  import doctest

  doctest.testmod()





