

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

  '''simple docstring'''





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


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

      UpperCAmelCase    				=						[t[-1] for t in os.walk(os.path.join(a_ ,			os.listdir(a_  )[0] ,			'snapshots'  )  )]
      UpperCAmelCase    				=						[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     snake_case_      (  self  )				->					Optional[Any]:


    """simple docstring"""
    UpperCAmelCase				,   UpperCAmelCase    				=						FlaxStableDiffusionPipeline.from_pretrained(
        'hf-internal-testing/tiny-stable-diffusion-pipe' ,			safety_checker=a_  )

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

    UpperCAmelCase    				=						jax.random.PRNGKey(0  )
    UpperCAmelCase    				=						4

    UpperCAmelCase    				=						jax.device_count()
    UpperCAmelCase    				=						num_samples * [prompt]
    UpperCAmelCase    				=						pipeline.prepare_inputs(a_  )

    # shard inputs and rng
    UpperCAmelCase    				=						replicate(a_  )
    UpperCAmelCase    				=						jax.random.split(a_ ,			a_  )
    UpperCAmelCase    				=						shard(a_  )

    UpperCAmelCase    				=						pipeline(a_ ,			a_ ,			a_ ,			a_ ,			jit=a_  ).images

    assert images.shape == (num_samples, 1, 6_4, 6_4, 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(a_ ,			dtype=np.floataa  ).sum() - 4_9947.875  ) < 5E-1

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


    """simple docstring"""
    UpperCAmelCase				,   UpperCAmelCase    				=						FlaxStableDiffusionPipeline.from_pretrained(
        'CompVis/stable-diffusion-v1-4' ,			revision='flax' ,			safety_checker=a_  )

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

    UpperCAmelCase    				=						jax.random.PRNGKey(0  )
    UpperCAmelCase    				=						5_0

    UpperCAmelCase    				=						jax.device_count()
    UpperCAmelCase    				=						num_samples * [prompt]
    UpperCAmelCase    				=						pipeline.prepare_inputs(a_  )

    # shard inputs and rng
    UpperCAmelCase    				=						replicate(a_  )
    UpperCAmelCase    				=						jax.random.split(a_ ,			a_  )
    UpperCAmelCase    				=						shard(a_  )

    UpperCAmelCase    				=						pipeline(a_ ,			a_ ,			a_ ,			a_ ,			jit=a_  ).images

    assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(a_ ,			dtype=np.floataa  ).sum() - 238_3808.2)  ) < 5E-1
  def     snake_case_      (  self  )				->					int:


    """simple docstring"""
    UpperCAmelCase				,   UpperCAmelCase    				=						FlaxStableDiffusionPipeline.from_pretrained(
        'CompVis/stable-diffusion-v1-4' ,			revision='bf16' ,			dtype=jnp.bfloataa ,			safety_checker=a_  )

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

    UpperCAmelCase    				=						jax.random.PRNGKey(0  )
    UpperCAmelCase    				=						5_0

    UpperCAmelCase    				=						jax.device_count()
    UpperCAmelCase    				=						num_samples * [prompt]
    UpperCAmelCase    				=						pipeline.prepare_inputs(a_  )

    # shard inputs and rng
    UpperCAmelCase    				=						replicate(a_  )
    UpperCAmelCase    				=						jax.random.split(a_ ,			a_  )
    UpperCAmelCase    				=						shard(a_  )

    UpperCAmelCase    				=						pipeline(a_ ,			a_ ,			a_ ,			a_ ,			jit=a_  ).images

    assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(a_ ,			dtype=np.floataa  ).sum() - 237_3516.75)  ) < 5E-1
  def     snake_case_      (  self  )				->					int:


    """simple docstring"""
    UpperCAmelCase				,   UpperCAmelCase    				=						FlaxStableDiffusionPipeline.from_pretrained(
        'CompVis/stable-diffusion-v1-4' ,			revision='bf16' ,			dtype=jnp.bfloataa  )

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

    UpperCAmelCase    				=						jax.random.PRNGKey(0  )
    UpperCAmelCase    				=						5_0

    UpperCAmelCase    				=						jax.device_count()
    UpperCAmelCase    				=						num_samples * [prompt]
    UpperCAmelCase    				=						pipeline.prepare_inputs(a_  )

    # shard inputs and rng
    UpperCAmelCase    				=						replicate(a_  )
    UpperCAmelCase    				=						jax.random.split(a_ ,			a_  )
    UpperCAmelCase    				=						shard(a_  )

    UpperCAmelCase    				=						pipeline(a_ ,			a_ ,			a_ ,			a_ ,			jit=a_  ).images

    assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(a_ ,			dtype=np.floataa  ).sum() - 237_3516.75)  ) < 5E-1
  def     snake_case_      (  self  )				->					Union[str, Any]:


    """simple docstring"""
    UpperCAmelCase    				=						FlaxDDIMScheduler(
        beta_start=0.0_0085 ,			beta_end=0.012 ,			beta_schedule='scaled_linear' ,			set_alpha_to_one=a_ ,			steps_offset=1 ,			)

    UpperCAmelCase				,   UpperCAmelCase    				=						FlaxStableDiffusionPipeline.from_pretrained(
        'CompVis/stable-diffusion-v1-4' ,			revision='bf16' ,			dtype=jnp.bfloataa ,			scheduler=a_ ,			safety_checker=a_ ,			)
    UpperCAmelCase    				=						scheduler.create_state()

    UpperCAmelCase    				=						scheduler_state

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

    UpperCAmelCase    				=						jax.random.PRNGKey(0  )
    UpperCAmelCase    				=						5_0

    UpperCAmelCase    				=						jax.device_count()
    UpperCAmelCase    				=						num_samples * [prompt]
    UpperCAmelCase    				=						pipeline.prepare_inputs(a_  )

    # shard inputs and rng
    UpperCAmelCase    				=						replicate(a_  )
    UpperCAmelCase    				=						jax.random.split(a_ ,			a_  )
    UpperCAmelCase    				=						shard(a_  )

    UpperCAmelCase    				=						pipeline(a_ ,			a_ ,			a_ ,			a_ ,			jit=a_  ).images

    assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 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(a_ ,			dtype=np.floataa  ).sum() - 234_7693.5)  ) < 5E-1





  def     snake_case_      (  self  )				->					Optional[Any]:


    """simple docstring"""
    UpperCAmelCase    				=						(
        'A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of'
        ' field, close up, split lighting, cinematic'
    )

    UpperCAmelCase    				=						jax.device_count()
    UpperCAmelCase    				=						num_samples * [prompt]
    UpperCAmelCase    				=						jax.random.split(jax.random.PRNGKey(0  ) ,			a_  )

    UpperCAmelCase				,   UpperCAmelCase    				=						FlaxStableDiffusionPipeline.from_pretrained(
        'CompVis/stable-diffusion-v1-4' ,			revision='bf16' ,			dtype=jnp.bfloataa ,			safety_checker=a_ ,			)

    UpperCAmelCase    				=						replicate(a_  )
    UpperCAmelCase    				=						pipeline.prepare_inputs(a_  )
    UpperCAmelCase    				=						shard(a_  )
    UpperCAmelCase    				=						pipeline(a_ ,			a_ ,			a_ ,			jit=a_  ).images
    assert images.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
    UpperCAmelCase    				=						images[2, 0, 2_5_6, 1_0:1_7, 1]

    # With memory efficient attention
    UpperCAmelCase				,   UpperCAmelCase    				=						FlaxStableDiffusionPipeline.from_pretrained(
        'CompVis/stable-diffusion-v1-4' ,			revision='bf16' ,			dtype=jnp.bfloataa ,			safety_checker=a_ ,			use_memory_efficient_attention=a_ ,			)

    UpperCAmelCase    				=						replicate(a_  )
    UpperCAmelCase    				=						pipeline.prepare_inputs(a_  )
    UpperCAmelCase    				=						shard(a_  )
    UpperCAmelCase    				=						pipeline(a_ ,			a_ ,			a_ ,			jit=a_  ).images
    assert images_eff.shape == (num_samples, 1, 5_1_2, 5_1_2, 3)
    UpperCAmelCase    				=						images[2, 0, 2_5_6, 1_0:1_7, 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 ..utils import DummyObject, requires_backends



class 							lowercase_						(		metaclass=a			):

  '''simple docstring'''





  __lowerCAmelCase       :			str     	= ["keras_nlp"]
  def __init__(  self ,			*a_ ,			**a_  )				->					Union[str, Any]:


    """simple docstring"""
    requires_backends(self ,			['keras_nlp']  )


