


'''simple docstring'''






__A					:  List[str]											=      [
    "Audio",
    "Array2D",
    "Array3D",
    "Array4D",
    "Array5D",
    "ClassLabel",
    "Features",
    "Sequence",
    "Value",
    "Image",
    "Translation",
    "TranslationVariableLanguages",
]
from .audio import Audio
from .features import ArrayaD, ArrayaD, ArrayaD, ArrayaD, ClassLabel, Features, Sequence, Value
from .image import Image
from .translation import Translation, TranslationVariableLanguages







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







    """simple docstring"""



    def __lowercase   (   self    :  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=lowerCamelCase    ,						cache_dir=lowerCamelCase		)

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







    """simple docstring"""



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

      lowerCAmelCase_			:      List[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_			:      Tuple			=   jax.random.PRNGKey(0		)
      lowerCAmelCase_			:      Optional[Any]			=   4

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

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

      lowerCAmelCase_			:      str			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).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_514_745		) < 1E-3
        assert np.abs(np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 49_947.875		) < 5E-1

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

    def __lowercase   (   self    :  Optional[int]		)   ->  List[str]:
      lowerCAmelCase_,							lowerCAmelCase_			:      Tuple			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""flax"""    ,						safety_checker=lowerCamelCase		)

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

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

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

      lowerCAmelCase_			:      Optional[int]			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      if jax.device_count() == 8:
        assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    ,						dtype=np.floataa		).sum() - 0.05_652_401)		) < 1E-3
        assert np.abs((np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 2_383_808.2)		) < 5E-1

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

      lowerCAmelCase_			:      Optional[int]			=   (
          """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_			:      Dict			=   jax.device_count()
      lowerCAmelCase_			:      List[Any]			=   num_samples * [prompt]
      lowerCAmelCase_			:      Tuple			=   pipeline.prepare_inputs(lowerCamelCase		)

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

      lowerCAmelCase_			:      Any			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      if jax.device_count() == 8:
        assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    ,						dtype=np.floataa		).sum() - 0.04_003_906)		) < 1E-3
        assert np.abs((np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 2_373_516.75)		) < 5E-1

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

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

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

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

      lowerCAmelCase_			:      Tuple			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      if jax.device_count() == 8:
        assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    ,						dtype=np.floataa		).sum() - 0.04_003_906)		) < 1E-3
        assert np.abs((np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 2_373_516.75)		) < 5E-1

    def __lowercase   (   self    :  List[str]		)   ->  Any:
      lowerCAmelCase_			:      List[Any]			=   FlaxDDIMScheduler(
          beta_start=0.00_085    ,						beta_end=0.012    ,						beta_schedule="""scaled_linear"""    ,						set_alpha_to_one=lowerCamelCase    ,						steps_offset=1    ,						)

      lowerCAmelCase_,							lowerCAmelCase_			:      List[str]			=   FlaxStableDiffusionPipeline.from_pretrained(
          """CompVis/stable-diffusion-v1-4"""    ,						revision="""bf16"""    ,						dtype=jnp.bfloataa    ,						scheduler=lowerCamelCase    ,						safety_checker=lowerCamelCase    ,						)
      lowerCAmelCase_			:      Any			=   scheduler.create_state()

      lowerCAmelCase_			:      Optional[Any]			=   scheduler_state

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

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

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

      lowerCAmelCase_			:      Tuple			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images

      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      if jax.device_count() == 8:
        assert np.abs((np.abs(images[0, 0, :2, :2, -2:]    ,						dtype=np.floataa		).sum() - 0.045_043_945)		) < 1E-3
        assert np.abs((np.abs(lowerCamelCase    ,						dtype=np.floataa		).sum() - 2_347_693.5)		) < 5E-1





    def __lowercase   (   self    :  Any		)   ->  Union[str, Any]:
      lowerCAmelCase_			:      int			=   (
          """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.device_count()
      lowerCAmelCase_			:      Optional[Any]			=   num_samples * [prompt]
      lowerCAmelCase_			:      Union[str, Any]			=   jax.random.split(jax.random.PRNGKey(0		)    ,						lowerCamelCase		)

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

      lowerCAmelCase_			:      Optional[int]			=   replicate(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   pipeline.prepare_inputs(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   shard(lowerCamelCase		)
      lowerCAmelCase_			:      Tuple			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images
      assert images.shape == (num_samples, 1, 5_12, 5_12, 3)
      lowerCAmelCase_			:      Dict			=   images[2, 0, 2_56, 10:17, 1]

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

      lowerCAmelCase_			:      List[str]			=   replicate(lowerCamelCase		)
      lowerCAmelCase_			:      Dict			=   pipeline.prepare_inputs(lowerCamelCase		)
      lowerCAmelCase_			:      Union[str, Any]			=   shard(lowerCamelCase		)
      lowerCAmelCase_			:      Optional[int]			=   pipeline(lowerCamelCase    ,						lowerCamelCase    ,						lowerCamelCase    ,						jit=lowerCamelCase		).images
      assert images_eff.shape == (num_samples, 1, 5_12, 5_12, 3)
      lowerCAmelCase_			:      List[Any]			=   images[2, 0, 2_56, 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



