





'''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   A_			(   self						:   Optional[int]			)    ->			Optional[int]:





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

         lowerCamelCase__					:   Union[str, Any]      				=	[t[-1] for t in os.walk(os.path.join(UpperCamelCase__    ,  os.listdir(UpperCamelCase__			)[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      lowerCAmelCase							(      unittest.TestCase    ):
 def   A_			(   self						:   Any			)    ->			List[str]:
     lowerCamelCase__					:   List[Any]      				=	FlaxStableDiffusionPipeline.from_pretrained(
         'hf-internal-testing/tiny-stable-diffusion-pipe'    ,  safety_checker=UpperCamelCase__			)

     lowerCamelCase__					:   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__					:   Any      				=	jax.random.PRNGKey(0			)
     lowerCamelCase__					:   Optional[int]      				=	4

     lowerCamelCase__					:   Optional[int]      				=	jax.device_count()
     lowerCamelCase__					:   Dict      				=	num_samples * [prompt]
     lowerCamelCase__					:   Any      				=	pipeline.prepare_inputs(UpperCamelCase__			)

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

     lowerCamelCase__					:   List[str]      				=	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.1_5_1_4_7_4_5			) < 1e-3
         assert np.abs(np.abs(UpperCamelCase__    ,  dtype=np.floataa			).sum() - 49947.875			) < 5e-1

     lowerCamelCase__					:   str      				=	pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]			)			)			)
     assert len(UpperCamelCase__			) == num_samples
 def   A_			(   self						:   Optional[Any]			)    ->			str:
     lowerCamelCase__					:   Optional[Any]      				=	FlaxStableDiffusionPipeline.from_pretrained(
         'CompVis/stable-diffusion-v1-4'    ,  revision='flax'    ,  safety_checker=UpperCamelCase__			)

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

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

     # shard inputs and rng
     lowerCamelCase__					:   Optional[Any]      				=	replicate(UpperCamelCase__			)
     lowerCamelCase__					:   Dict      				=	jax.random.split(UpperCamelCase__    ,  UpperCamelCase__			)
     lowerCamelCase__					:   Union[str, Any]      				=	shard(UpperCamelCase__			)

     lowerCamelCase__					:   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.0_5_6_5_2_4_0_1)			) < 1e-3
         assert np.abs((np.abs(UpperCamelCase__    ,  dtype=np.floataa			).sum() - 2383808.2)			) < 5e-1
 def   A_			(   self						:   Optional[Any]			)    ->			List[str]:
     lowerCamelCase__					:   Union[str, Any]      				=	FlaxStableDiffusionPipeline.from_pretrained(
         'CompVis/stable-diffusion-v1-4'    ,  revision='bf16'    ,  dtype=jnp.bfloataa    ,  safety_checker=UpperCamelCase__			)

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

     lowerCamelCase__					:   Optional[int]      				=	jax.device_count()
     lowerCamelCase__					:   int      				=	num_samples * [prompt]
     lowerCamelCase__					:   List[Any]      				=	pipeline.prepare_inputs(UpperCamelCase__			)

     # shard inputs and rng
     lowerCamelCase__					:   str      				=	replicate(UpperCamelCase__			)
     lowerCamelCase__					:   Optional[Any]      				=	jax.random.split(UpperCamelCase__    ,  UpperCamelCase__			)
     lowerCamelCase__					:   Dict      				=	shard(UpperCamelCase__			)

     lowerCamelCase__					:   Tuple      				=	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.0_4_0_0_3_9_0_6)			) < 1e-3
         assert np.abs((np.abs(UpperCamelCase__    ,  dtype=np.floataa			).sum() - 2373516.75)			) < 5e-1
 def   A_			(   self						:   Dict			)    ->			Union[str, Any]:
     lowerCamelCase__					:   Dict      				=	FlaxStableDiffusionPipeline.from_pretrained(
         'CompVis/stable-diffusion-v1-4'    ,  revision='bf16'    ,  dtype=jnp.bfloataa			)

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

     lowerCamelCase__					:   Optional[int]      				=	jax.device_count()
     lowerCamelCase__					:   Union[str, Any]      				=	num_samples * [prompt]
     lowerCamelCase__					:   str      				=	pipeline.prepare_inputs(UpperCamelCase__			)

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

     lowerCamelCase__					:   Optional[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.0_4_0_0_3_9_0_6)			) < 1e-3
         assert np.abs((np.abs(UpperCamelCase__    ,  dtype=np.floataa			).sum() - 2373516.75)			) < 5e-1
 def   A_			(   self						:   int			)    ->			Dict:
     lowerCamelCase__					:   Tuple      				=	FlaxDDIMScheduler(
         beta_start=0.0_0_0_8_5    ,  beta_end=0.0_1_2    ,  beta_schedule='scaled_linear'    ,  set_alpha_to_one=UpperCamelCase__    ,  steps_offset=1    ,  )

     lowerCamelCase__					:   Any      				=	FlaxStableDiffusionPipeline.from_pretrained(
         'CompVis/stable-diffusion-v1-4'    ,  revision='bf16'    ,  dtype=jnp.bfloataa    ,  scheduler=UpperCamelCase__    ,  safety_checker=UpperCamelCase__    ,  )
     lowerCamelCase__					:   Any      				=	scheduler.create_state()

     lowerCamelCase__					:   Tuple      				=	scheduler_state

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

     lowerCamelCase__					:   Tuple      				=	jax.device_count()
     lowerCamelCase__					:   Optional[Any]      				=	num_samples * [prompt]
     lowerCamelCase__					:   str      				=	pipeline.prepare_inputs(UpperCamelCase__			)

     # shard inputs and rng
     lowerCamelCase__					:   Any      				=	replicate(UpperCamelCase__			)
     lowerCamelCase__					:   Optional[Any]      				=	jax.random.split(UpperCamelCase__    ,  UpperCamelCase__			)
     lowerCamelCase__					:   Dict      				=	shard(UpperCamelCase__			)

     lowerCamelCase__					:   Optional[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.0_4_5_0_4_3_9_4_5)			) < 1e-3
         assert np.abs((np.abs(UpperCamelCase__    ,  dtype=np.floataa			).sum() - 2347693.5)			) < 5e-1







 def   A_			(   self						:   Optional[int]			)    ->			int:
     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__					:   int      				=	jax.device_count()
     lowerCamelCase__					:   Any      				=	num_samples * [prompt]
     lowerCamelCase__					:   List[str]      				=	jax.random.split(jax.random.PRNGKey(0			)    ,  UpperCamelCase__			)

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

     lowerCamelCase__					:   str      				=	replicate(UpperCamelCase__			)
     lowerCamelCase__					:   List[str]      				=	pipeline.prepare_inputs(UpperCamelCase__			)
     lowerCamelCase__					:   Dict      				=	shard(UpperCamelCase__			)
     lowerCamelCase__					:   Tuple      				=	pipeline(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__    ,  jit=UpperCamelCase__			).images
     assert images.shape == (num_samples, 1, 512, 512, 3)
     lowerCamelCase__					:   Any      				=	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=UpperCamelCase__    ,  use_memory_efficient_attention=UpperCamelCase__    ,  )

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












import enum
import shutil
import sys


_UpperCAmelCase     ,_UpperCAmelCase       : List[Any]			     =       shutil.get_terminal_size()

_UpperCAmelCase       : Union[str, Any]			     =       {"""UP""": """A""", """DOWN""": """B""", """RIGHT""": """C""", """LEFT""": """D"""}

class      lowerCAmelCase							(      enum.Enum    ):
 UpperCAmelCase__									=	0
 UpperCAmelCase__									=	1


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase=""				)    ->  int:
    sys.stdout.write(str(_UpperCAmelCase				) + end				)
    sys.stdout.flush()


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase		,			_UpperCAmelCase=""				)    ->  Union[str, Any]:
    forceWrite(F"""\u001b[{color}m{content}\u001b[0m"""		,			_UpperCAmelCase				)


def  SCREAMING_SNAKE_CASE   ( )    ->  Dict:
    forceWrite('\r'				)


def  SCREAMING_SNAKE_CASE   ( _UpperCAmelCase		,			_UpperCAmelCase				)    ->  Optional[Any]:
    forceWrite(F"""\033[{num_lines}{CURSOR_TO_CHAR[direction.upper()]}"""				)


def  SCREAMING_SNAKE_CASE   ( )    ->  int:
    forceWrite(' ' * TERMINAL_WIDTH				)
    reset_cursor()


def  SCREAMING_SNAKE_CASE   ( )    ->  Optional[Any]:
    reset_cursor()
    forceWrite('-' * TERMINAL_WIDTH				)





