





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

   '''simple docstring'''


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






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

           UpperCamelCase__							:				Union[str, Any]          =      [t[-1] for t in os.walk(os.path.join(__magic_name__,		os.listdir(__magic_name__							)[0],		'''snapshots'''							)							)]
           UpperCamelCase__							:				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     lowercase__       (	unittest.TestCase				):

   '''simple docstring'''


   def       UpperCamelCase__  (    self							) ->      Optional[int]:






       """simple docstring"""
       UpperCamelCase__   ,UpperCamelCase__							:				List[Any]          =      FlaxStableDiffusionPipeline.from_pretrained(
           '''hf-internal-testing/tiny-stable-diffusion-pipe''',		safety_checker=__magic_name__							)

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

       UpperCamelCase__							:				Tuple          =      jax.random.PRNGKey(0							)
       UpperCamelCase__							:				str          =      4

       UpperCamelCase__							:				Dict          =      jax.device_count()
       UpperCamelCase__							:				Dict          =      num_samples * [prompt]
       UpperCamelCase__							:				Dict          =      pipeline.prepare_inputs(__magic_name__							)

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

       UpperCamelCase__							:				str          =      pipeline(__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		jit=__magic_name__							).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.151_4745							) < 1E-3
           assert np.abs(np.abs(__magic_name__,		dtype=np.floataa							).sum() - 4_9947.875							) < 5E-1

       UpperCamelCase__							:				Union[str, Any]          =      pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]							)							)							)
       assert len(__magic_name__							) == num_samples




   def       UpperCamelCase__  (    self							) ->      Optional[int]:






       """simple docstring"""
       UpperCamelCase__   ,UpperCamelCase__							:				Optional[Any]          =      FlaxStableDiffusionPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4''',		revision='''flax''',		safety_checker=__magic_name__							)

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

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

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

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

       UpperCamelCase__							:				Optional[int]          =      pipeline(__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		jit=__magic_name__							).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.0565_2401)							) < 1E-3
           assert np.abs((np.abs(__magic_name__,		dtype=np.floataa							).sum() - 238_3808.2)							) < 5E-1




   def       UpperCamelCase__  (    self							) ->      str:






       """simple docstring"""
       UpperCamelCase__   ,UpperCamelCase__							:				Optional[int]          =      FlaxStableDiffusionPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4''',		revision='''bf16''',		dtype=jnp.bfloataa,		safety_checker=__magic_name__							)

       UpperCamelCase__							:				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'''
       )

       UpperCamelCase__							:				Tuple          =      jax.random.PRNGKey(0							)
       UpperCamelCase__							:				Any          =      50

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

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

       UpperCamelCase__							:				Union[str, Any]          =      pipeline(__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		jit=__magic_name__							).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.0400_3906)							) < 1E-3
           assert np.abs((np.abs(__magic_name__,		dtype=np.floataa							).sum() - 237_3516.75)							) < 5E-1




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






       """simple docstring"""
       UpperCamelCase__   ,UpperCamelCase__							:				Optional[int]          =      FlaxStableDiffusionPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4''',		revision='''bf16''',		dtype=jnp.bfloataa							)

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

       UpperCamelCase__							:				int          =      jax.random.PRNGKey(0							)
       UpperCamelCase__							:				Tuple          =      50

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

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

       UpperCamelCase__							:				Optional[Any]          =      pipeline(__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		jit=__magic_name__							).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.0400_3906)							) < 1E-3
           assert np.abs((np.abs(__magic_name__,		dtype=np.floataa							).sum() - 237_3516.75)							) < 5E-1




   def       UpperCamelCase__  (    self							) ->      Tuple:






       """simple docstring"""
       UpperCamelCase__							:				Any          =      FlaxDDIMScheduler(
           beta_start=0.0_0085,		beta_end=0.012,		beta_schedule='''scaled_linear''',		set_alpha_to_one=__magic_name__,		steps_offset=1,		)

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

       UpperCamelCase__							:				Union[str, Any]          =      scheduler_state

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

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

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

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

       UpperCamelCase__							:				List[Any]          =      pipeline(__magic_name__,		__magic_name__,		__magic_name__,		__magic_name__,		jit=__magic_name__							).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_4504_3945)							) < 1E-3
           assert np.abs((np.abs(__magic_name__,		dtype=np.floataa							).sum() - 234_7693.5)							) < 5E-1




   def       UpperCamelCase__  (    self							) ->      Any:






       """simple docstring"""
       UpperCamelCase__							:				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'''
       )

       UpperCamelCase__							:				Optional[int]          =      jax.device_count()
       UpperCamelCase__							:				Tuple          =      num_samples * [prompt]
       UpperCamelCase__							:				Optional[Any]          =      jax.random.split(jax.random.PRNGKey(0							),		__magic_name__							)

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

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

       # With memory efficient attention
       UpperCamelCase__   ,UpperCamelCase__							:				List[str]          =      FlaxStableDiffusionPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4''',		revision='''bf16''',		dtype=jnp.bfloataa,		safety_checker=__magic_name__,		use_memory_efficient_attention=__magic_name__,		)

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




def 		lowerCAmelCase_    (  __UpperCAmelCase:			list		,   __UpperCAmelCase:			Optional[int]      )      ->	tuple:
    UpperCamelCase__   ,UpperCamelCase__   ,UpperCamelCase__							:				List[Any]          =      [], [], []
    for element in data:
        if element < pivot:
            less.append(__UpperCAmelCase      )
        elif element > pivot:
            greater.append(__UpperCAmelCase      )
        else:
            equal.append(__UpperCAmelCase      )
    return less, equal, greater





def 		lowerCAmelCase_    (  __UpperCAmelCase:			list		,   __UpperCAmelCase:			int      )      ->	List[str]:
    # index = len(items) // 2 when trying to find the median
    #   (value of index when items is sorted)

    # invalid input
    if index >= len(__UpperCAmelCase      ) or index < 0:
        return None

    UpperCamelCase__							:				List[str]          =      items[random.randint(0		,   len(__UpperCAmelCase      ) - 1      )]
    UpperCamelCase__							:				List[Any]          =      0
    UpperCamelCase__   ,UpperCamelCase__   ,UpperCamelCase__							:				int          =      _partition(__UpperCAmelCase		,   __UpperCAmelCase      )
    UpperCamelCase__							:				Union[str, Any]          =      len(__UpperCAmelCase      )
    UpperCamelCase__							:				Dict          =      len(__UpperCAmelCase      )

    # index is the pivot
    if m <= index < m + count:
        return pivot
    # must be in smaller
    elif m > index:
        return quick_select(__UpperCAmelCase		,   __UpperCAmelCase      )
    # must be in larger
    else:
        return quick_select(__UpperCAmelCase		,   index - (m + count)      )

