




'''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 							lowercase_						( self   :  List[str]       ):


           '''simple docstring'''

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

                a_     :							Dict     						=  [t[-1] for t in os.walk(os.path.join(lowercase__						,						os.listdir(lowercase__       )[0]						,						"""snapshots"""       )       )]
                a_     :							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      SCREAMING_SNAKE_CASE    (						unittest.TestCase				):






      def 							lowercase_						( self   :  Tuple       ):


           '''simple docstring'''

           a_						,     a_     :							Optional[int]     						=  FlaxStableDiffusionPipeline.from_pretrained(
               """hf-internal-testing/tiny-stable-diffusion-pipe"""						,						safety_checker=lowercase__       )

           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_     :							Optional[Any]     						=  jax.random.PRNGKey(0       )
           a_     :							Union[str, Any]     						=  4

           a_     :							List[Any]     						=  jax.device_count()
           a_     :							Dict     						=  num_samples * [prompt]
           a_     :							Optional[Any]     						=  pipeline.prepare_inputs(lowercase__       )

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

           a_     :							Tuple     						=  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() - 49947.875       ) < 5e-1

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






      def 							lowercase_						( self   :  str       ):


           '''simple docstring'''

           a_						,     a_     :							Any     						=  FlaxStableDiffusionPipeline.from_pretrained(
               """CompVis/stable-diffusion-v1-4"""						,						revision="""flax"""						,						safety_checker=lowercase__       )

           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_     :							Dict     						=  jax.random.PRNGKey(0       )
           a_     :							Tuple     						=  50

           a_     :							Union[str, Any]     						=  jax.device_count()
           a_     :							Dict     						=  num_samples * [prompt]
           a_     :							Optional[Any]     						=  pipeline.prepare_inputs(lowercase__       )

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

           a_     :							Optional[int]     						=  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() - 2383808.2)       ) < 5e-1






      def 							lowercase_						( self   :  Optional[Any]       ):


           '''simple docstring'''

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

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

           a_     :							str     						=  jax.random.PRNGKey(0       )
           a_     :							Tuple     						=  50

           a_     :							List[Any]     						=  jax.device_count()
           a_     :							Optional[Any]     						=  num_samples * [prompt]
           a_     :							Tuple     						=  pipeline.prepare_inputs(lowercase__       )

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

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






      def 							lowercase_						( self   :  Optional[Any]       ):


           '''simple docstring'''

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

           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_     :							Union[str, Any]     						=  jax.random.PRNGKey(0       )
           a_     :							List[str]     						=  50

           a_     :							Any     						=  jax.device_count()
           a_     :							Tuple     						=  num_samples * [prompt]
           a_     :							Optional[Any]     						=  pipeline.prepare_inputs(lowercase__       )

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

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






      def 							lowercase_						( self   :  Any       ):


           '''simple docstring'''

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

           a_						,     a_     :							Tuple     						=  FlaxStableDiffusionPipeline.from_pretrained(
               """CompVis/stable-diffusion-v1-4"""						,						revision="""bf16"""						,						dtype=jnp.bfloataa						,						scheduler=lowercase__						,						safety_checker=lowercase__						,						)
           a_     :							Optional[int]     						=  scheduler.create_state()

           a_     :							Dict     						=  scheduler_state

           a_     :							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[str]     						=  50

           a_     :							int     						=  jax.device_count()
           a_     :							Any     						=  num_samples * [prompt]
           a_     :							Any     						=  pipeline.prepare_inputs(lowercase__       )

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

           a_     :							int     						=  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() - 2347693.5)       ) < 5e-1






      def 							lowercase_						( self   :  int       ):


           '''simple docstring'''

           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.device_count()
           a_     :							str     						=  num_samples * [prompt]
           a_     :							int     						=  jax.random.split(jax.random.PRNGKey(0       )						,						lowercase__       )

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

           a_     :							Dict     						=  replicate(lowercase__       )
           a_     :							Optional[Any]     						=  pipeline.prepare_inputs(lowercase__       )
           a_     :							List[Any]     						=  shard(lowercase__       )
           a_     :							Dict     						=  pipeline(lowercase__						,						lowercase__						,						lowercase__						,						jit=lowercase__       ).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=lowercase__						,						use_memory_efficient_attention=lowercase__						,						)

           a_     :							str     						=  replicate(lowercase__       )
           a_     :							Any     						=  pipeline.prepare_inputs(lowercase__       )
           a_     :							Union[str, Any]     						=  shard(lowercase__       )
           a_     :							List[str]     						=  pipeline(lowercase__						,						lowercase__						,						lowercase__						,						jit=lowercase__       ).images
           assert images_eff.shape == (num_samples, 1, 512, 512, 3)
           a_     :							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'''


from __future__ import annotations



def  _SCREAMING_SNAKE_CASE     (      UpperCamelCase__    :					list[float]				):




     """simple docstring"""


     if len(UpperCamelCase__				) < 2:
          raise ValueError("""Monogons and Digons are not polygons in the Euclidean space"""				)
     if any(i <= 0 for i in nums				):
          raise ValueError("""All values must be greater than 0"""				)
     a_     :							int     						=  nums.copy()
     copy_nums.sort()
     return copy_nums[-1] < sum(copy_nums[:-1]				)


if __name__ == "__main__":
 import doctest

 doctest.testmod()
