





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







  '''simple docstring'''


  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					List[Any]:
        with tempfile.TemporaryDirectory() as tmpdirname:
              # pipeline has Flax weights
              _UpperCamelCase   :						Optional[int]    =							FlaxDiffusionPipeline.from_pretrained(
                  "hf-internal-testing/tiny-stable-diffusion-pipe"       , safety_checker=__a       , cache_dir=__a						)

              _UpperCamelCase   :						int    =							[t[-1] for t in os.walk(os.path.join(__a       , os.listdir(__a						)[0]       , "snapshots"						)						)]
              _UpperCamelCase   :						Optional[int]    =							[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   __SCREAMING_SNAKE_CASE     (    unittest.TestCase ):







  '''simple docstring'''


  def  __SCREAMING_SNAKE_CASE					(		self :							List[str]						)		->					Any:
        _UpperCamelCase,	_UpperCamelCase   :						Dict    =							FlaxStableDiffusionPipeline.from_pretrained(
            "hf-internal-testing/tiny-stable-diffusion-pipe"       , safety_checker=__a						)

        _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   :						int    =							jax.random.PRNGKey(0						)
        _UpperCamelCase   :						int    =							4

        _UpperCamelCase   :						List[str]    =							jax.device_count()
        _UpperCamelCase   :						str    =							num_samples * [prompt]
        _UpperCamelCase   :						str    =							pipeline.prepare_inputs(__a						)

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

        _UpperCamelCase   :						Dict    =							pipeline(__a       , __a       , __a       , __a       , jit=__a						).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(__a       , dtype=np.floataa						).sum() - 4_99_47.8_75						) < 5e-1

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

  def  __SCREAMING_SNAKE_CASE					(		self :							Union[str, Any]						)		->					Optional[int]:
        _UpperCamelCase,	_UpperCamelCase   :						Tuple    =							FlaxStableDiffusionPipeline.from_pretrained(
            "CompVis/stable-diffusion-v1-4"       , revision="flax"       , safety_checker=__a						)

        _UpperCamelCase   :						str    =							(
            "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   :						Union[str, Any]    =							50

        _UpperCamelCase   :						Union[str, Any]    =							jax.device_count()
        _UpperCamelCase   :						Optional[Any]    =							num_samples * [prompt]
        _UpperCamelCase   :						Tuple    =							pipeline.prepare_inputs(__a						)

        # shard inputs and rng
        _UpperCamelCase   :						str    =							replicate(__a						)
        _UpperCamelCase   :						Optional[int]    =							jax.random.split(__a       , __a						)
        _UpperCamelCase   :						int    =							shard(__a						)

        _UpperCamelCase   :						Union[str, Any]    =							pipeline(__a       , __a       , __a       , __a       , jit=__a						).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(__a       , dtype=np.floataa						).sum() - 2_38_38_08.2)						) < 5e-1

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

        _UpperCamelCase   :						int    =							(
            "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[Any]    =							num_samples * [prompt]
        _UpperCamelCase   :						Tuple    =							pipeline.prepare_inputs(__a						)

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

        _UpperCamelCase   :						str    =							pipeline(__a       , __a       , __a       , __a       , jit=__a						).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(__a       , dtype=np.floataa						).sum() - 2_37_35_16.75)						) < 5e-1

  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					List[Any]:
        _UpperCamelCase,	_UpperCamelCase   :						List[Any]    =							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   :						Tuple    =							jax.random.PRNGKey(0						)
        _UpperCamelCase   :						Optional[Any]    =							50

        _UpperCamelCase   :						Any    =							jax.device_count()
        _UpperCamelCase   :						Any    =							num_samples * [prompt]
        _UpperCamelCase   :						Any    =							pipeline.prepare_inputs(__a						)

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

        _UpperCamelCase   :						int    =							pipeline(__a       , __a       , __a       , __a       , jit=__a						).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(__a       , dtype=np.floataa						).sum() - 2_37_35_16.75)						) < 5e-1

  def  __SCREAMING_SNAKE_CASE					(		self :							List[Any]						)		->					Optional[Any]:
        _UpperCamelCase   :						Optional[Any]    =							FlaxDDIMScheduler(
            beta_start=0.0_00_85       , beta_end=0.0_12       , beta_schedule="scaled_linear"       , set_alpha_to_one=__a       , steps_offset=1       , )

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

        _UpperCamelCase   :						str    =							scheduler_state

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

        _UpperCamelCase   :						List[Any]    =							jax.random.PRNGKey(0						)
        _UpperCamelCase   :						int    =							50

        _UpperCamelCase   :						int    =							jax.device_count()
        _UpperCamelCase   :						List[Any]    =							num_samples * [prompt]
        _UpperCamelCase   :						Dict    =							pipeline.prepare_inputs(__a						)

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

        _UpperCamelCase   :						List[str]    =							pipeline(__a       , __a       , __a       , __a       , jit=__a						).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(__a       , dtype=np.floataa						).sum() - 2_34_76_93.5)						) < 5e-1


  def  __SCREAMING_SNAKE_CASE					(		self :							Dict						)		->					Any:
        _UpperCamelCase   :						str    =							(
            "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
            " field, close up, split lighting, cinematic"
        )

        _UpperCamelCase   :						str    =							jax.device_count()
        _UpperCamelCase   :						Any    =							num_samples * [prompt]
        _UpperCamelCase   :						Dict    =							jax.random.split(jax.random.PRNGKey(0						)       , __a						)

        _UpperCamelCase,	_UpperCamelCase   :						Optional[int]    =							FlaxStableDiffusionPipeline.from_pretrained(
            "CompVis/stable-diffusion-v1-4"       , revision="bf16"       , dtype=jnp.bfloataa       , safety_checker=__a       , )

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

        _UpperCamelCase   :						Optional[Any]    =							replicate(__a						)
        _UpperCamelCase   :						Optional[int]    =							pipeline.prepare_inputs(__a						)
        _UpperCamelCase   :						Any    =							shard(__a						)
        _UpperCamelCase   :						Union[str, Any]    =							pipeline(__a       , __a       , __a       , jit=__a						).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







"""simple docstring"""


import random
from typing import Any




def        lowercase__     (       lowercase_      )       ->		list[Any]:




      """simple docstring"""

      for _ in range(len(lowercase_      )      ):
            _UpperCamelCase   :						Dict    =							random.randint(0		,len(lowercase_      ) - 1      )
            _UpperCamelCase   :						str    =							random.randint(0		,len(lowercase_      ) - 1      )
            _UpperCamelCase,	_UpperCamelCase   :						Optional[Any]    =							data[b], data[a]
      return data


if __name__ == "__main__":
  lowerCamelCase__						= [0, 1, 2, 3, 4, 5, 6, 7]
  lowerCamelCase__						= ["python", "says", "hello", "!"]
  print("Fisher-Yates Shuffle:")
  print("List", integers, strings)
  print("FY Shuffle", fisher_yates_shuffle(integers), fisher_yates_shuffle(strings))
