


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


     """simple docstring"""

     def 		a__    (						self:    List[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=lowercase__  ,		cache_dir=lowercase__						)

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


     """simple docstring"""

     def 		a__    (						self:    Tuple						)->     Optional[Any]:
           lowerCamelCase       :       Optional[int]   =						FlaxStableDiffusionPipeline.from_pretrained(
               """hf-internal-testing/tiny-stable-diffusion-pipe"""  ,		safety_checker=lowercase__						)

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

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

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

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

           lowerCamelCase       :       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.1_51_47_45						) < 1e-3
                 assert np.abs(np.abs(lowercase__  ,		dtype=np.floataa						).sum() - 4_99_47.8_75						) < 5e-1

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

     def 		a__    (						self:    str						)->     List[str]:
           lowerCamelCase       :       Any   =						FlaxStableDiffusionPipeline.from_pretrained(
               """CompVis/stable-diffusion-v1-4"""  ,		revision="""flax"""  ,		safety_checker=lowercase__						)

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

           lowerCamelCase       :       Dict   =						jax.random.PRNGKey(0						)
           lowerCamelCase       :       Tuple   =						50

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

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

           lowerCamelCase       :       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.05_65_24_01)						) < 1e-3
                 assert np.abs((np.abs(lowercase__  ,		dtype=np.floataa						).sum() - 2_38_38_08.2)						) < 5e-1

     def 		a__    (						self:    Optional[Any]						)->     int:
           lowerCamelCase       :       str   =						FlaxStableDiffusionPipeline.from_pretrained(
               """CompVis/stable-diffusion-v1-4"""  ,		revision="""bf16"""  ,		dtype=jnp.bfloataa  ,		safety_checker=lowercase__						)

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

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

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

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

           lowerCamelCase       :       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.04_00_39_06)						) < 1e-3
                 assert np.abs((np.abs(lowercase__  ,		dtype=np.floataa						).sum() - 2_37_35_16.75)						) < 5e-1

     def 		a__    (						self:    Optional[Any]						)->     Union[str, Any]:
           lowerCamelCase       :       int   =						FlaxStableDiffusionPipeline.from_pretrained(
               """CompVis/stable-diffusion-v1-4"""  ,		revision="""bf16"""  ,		dtype=jnp.bfloataa						)

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

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

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

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

           lowerCamelCase       :       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.04_00_39_06)						) < 1e-3
                 assert np.abs((np.abs(lowercase__  ,		dtype=np.floataa						).sum() - 2_37_35_16.75)						) < 5e-1

     def 		a__    (						self:    Any						)->     Any:
           lowerCamelCase       :       Optional[Any]   =						FlaxDDIMScheduler(
               beta_start=0.0_00_85  ,		beta_end=0.0_12  ,		beta_schedule="""scaled_linear"""  ,		set_alpha_to_one=lowercase__  ,		steps_offset=1  ,		)

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

           lowerCamelCase       :       Dict   =						scheduler_state

           lowerCamelCase       :       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       :       Optional[Any]   =						jax.random.PRNGKey(0						)
           lowerCamelCase       :       List[str]   =						50

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

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

           lowerCamelCase       :       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.0_45_04_39_45)						) < 1e-3
                 assert np.abs((np.abs(lowercase__  ,		dtype=np.floataa						).sum() - 2_34_76_93.5)						) < 5e-1







     def 		a__    (						self:    int						)->     List[Any]:
           lowerCamelCase       :       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"""
           )

           lowerCamelCase       :       Optional[int]   =						jax.device_count()
           lowerCamelCase       :       str   =						num_samples * [prompt]
           lowerCamelCase       :       int   =						jax.random.split(jax.random.PRNGKey(0						)  ,		lowercase__						)

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

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

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

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



def  _SCREAMING_SNAKE_CASE     (      UpperCamelCase__    :					float ,    UpperCamelCase__    :					bool = False				):




     """simple docstring"""


     if deriv:
          return value * (1 - value)
     return 1 / (1 + math.exp(-value				))


# Initial Value
lowerCAmelCase_      :  Optional[Any]				     =  0.02



def  _SCREAMING_SNAKE_CASE     (      UpperCamelCase__    :					int ,    UpperCamelCase__    :					int				):




     """simple docstring"""


     a_     :							Tuple     						=  float(2 * (random.randint(1 ,    100				)) - 1				)

     for _ in range(UpperCamelCase__				):
          # Forward propagation
          a_     :							Dict     						=  sigmoid_function(INITIAL_VALUE * weight				)
          # How much did we miss?
          a_     :							Dict     						=  (expected / 100) - layer_a
          # Error delta
          a_     :							Optional[Any]     						=  layer_1_error * sigmoid_function(UpperCamelCase__ ,    UpperCamelCase__				)
          # Update weight
          weight += INITIAL_VALUE * layer_1_delta

     return layer_a * 100


if __name__ == "__main__":
 import doctest

 doctest.testmod()

 lowerCAmelCase_      :  str				     =  int(input('Expected value: '))
 lowerCAmelCase_      :  str				     =  int(input('Number of propagations: '))
 print(forward_propagation(expected, number_propagations))
