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






     """simple docstring"""







     def       UpperCAmelCase__				(							self				:				Dict      ):







       """simple docstring"""




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

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






     """simple docstring"""







     def       UpperCAmelCase__				(							self				:				str      ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE	:     Any =						FlaxStableDiffusionPipeline.from_pretrained(
           '''hf-internal-testing/tiny-stable-diffusion-pipe'''			,		safety_checker=_A      )

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

       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						jax.random.PRNGKey(0      )
       __SCREAMING_SNAKE_CASE	:     Tuple =						4

       __SCREAMING_SNAKE_CASE	:     int =						jax.device_count()
       __SCREAMING_SNAKE_CASE	:     List[Any] =						num_samples * [prompt]
       __SCREAMING_SNAKE_CASE	:     int =						pipeline.prepare_inputs(_A      )

       # shard inputs and rng
       __SCREAMING_SNAKE_CASE	:     Any =						replicate(_A      )
       __SCREAMING_SNAKE_CASE	:     List[Any] =						jax.random.split(_A			,		_A      )
       __SCREAMING_SNAKE_CASE	:     List[Any] =						shard(_A      )

       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						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

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



     def       UpperCAmelCase__				(							self				:				Optional[Any]      ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE	:     Optional[Any] =						FlaxStableDiffusionPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4'''			,		revision='''flax'''			,		safety_checker=_A      )

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

       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						jax.random.PRNGKey(0      )
       __SCREAMING_SNAKE_CASE	:     Optional[int] =						50

       __SCREAMING_SNAKE_CASE	:     int =						jax.device_count()
       __SCREAMING_SNAKE_CASE	:     str =						num_samples * [prompt]
       __SCREAMING_SNAKE_CASE	:     str =						pipeline.prepare_inputs(_A      )

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

       __SCREAMING_SNAKE_CASE	:     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.05_65_24_01)      ) < 1e-3
         assert np.abs((np.abs(_A			,		dtype=np.floataa      ).sum() - 2_38_38_08.2)      ) < 5e-1



     def       UpperCAmelCase__				(							self				:				List[Any]      ):







       """simple docstring"""




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

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

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

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

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

       __SCREAMING_SNAKE_CASE	:     Tuple =						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       UpperCAmelCase__				(							self				:				Optional[int]      ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE	:     str =						FlaxStableDiffusionPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4'''			,		revision='''bf16'''			,		dtype=jnp.bfloataa      )

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

       __SCREAMING_SNAKE_CASE	:     Optional[Any] =						jax.random.PRNGKey(0      )
       __SCREAMING_SNAKE_CASE	:     Tuple =						50

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

       # shard inputs and rng
       __SCREAMING_SNAKE_CASE	:     Dict =						replicate(_A      )
       __SCREAMING_SNAKE_CASE	:     Dict =						jax.random.split(_A			,		_A      )
       __SCREAMING_SNAKE_CASE	:     Tuple =						shard(_A      )

       __SCREAMING_SNAKE_CASE	:     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.04_00_39_06)      ) < 1e-3
         assert np.abs((np.abs(_A			,		dtype=np.floataa      ).sum() - 2_37_35_16.75)      ) < 5e-1



     def       UpperCAmelCase__				(							self				:				Any      ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE	:     Dict =						FlaxDDIMScheduler(
           beta_start=0.0_00_85			,		beta_end=0.0_12			,		beta_schedule='''scaled_linear'''			,		set_alpha_to_one=_A			,		steps_offset=1			,		)

       __SCREAMING_SNAKE_CASE	:     Optional[Any] =						FlaxStableDiffusionPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4'''			,		revision='''bf16'''			,		dtype=jnp.bfloataa			,		scheduler=_A			,		safety_checker=_A			,		)
       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						scheduler.create_state()

       __SCREAMING_SNAKE_CASE	:     Optional[int] =						scheduler_state

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

       __SCREAMING_SNAKE_CASE	:     Optional[Any] =						jax.random.PRNGKey(0      )
       __SCREAMING_SNAKE_CASE	:     Optional[int] =						50

       __SCREAMING_SNAKE_CASE	:     List[str] =						jax.device_count()
       __SCREAMING_SNAKE_CASE	:     Dict =						num_samples * [prompt]
       __SCREAMING_SNAKE_CASE	:     List[str] =						pipeline.prepare_inputs(_A      )

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

       __SCREAMING_SNAKE_CASE	:     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       UpperCAmelCase__				(							self				:				str      ):







       """simple docstring"""




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

       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						jax.device_count()
       __SCREAMING_SNAKE_CASE	:     Any =						num_samples * [prompt]
       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						jax.random.split(jax.random.PRNGKey(0      )			,		_A      )

       __SCREAMING_SNAKE_CASE	:     Tuple =						FlaxStableDiffusionPipeline.from_pretrained(
           '''CompVis/stable-diffusion-v1-4'''			,		revision='''bf16'''			,		dtype=jnp.bfloataa			,		safety_checker=_A			,		)

       __SCREAMING_SNAKE_CASE	:     str =						replicate(_A      )
       __SCREAMING_SNAKE_CASE	:     Union[str, Any] =						pipeline.prepare_inputs(_A      )
       __SCREAMING_SNAKE_CASE	:     List[Any] =						shard(_A      )
       __SCREAMING_SNAKE_CASE	:     List[str] =						pipeline(_A			,		_A			,		_A			,		jit=_A      ).images
       assert images.shape == (num_samples, 1, 512, 512, 3)
       __SCREAMING_SNAKE_CASE	:     List[Any] =						images[2, 0, 256, 10:17, 1]

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

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

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


__lowerCamelCase					:	Dict       =  {'''configuration_vit_mae''': ['''VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ViTMAEConfig''']}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCamelCase					:	Any       =  [
		    '''VIT_MAE_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''ViTMAEForPreTraining''',
		    '''ViTMAELayer''',
		    '''ViTMAEModel''',
		    '''ViTMAEPreTrainedModel''',
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCamelCase					:	Optional[Any]       =  [
		    '''TFViTMAEForPreTraining''',
		    '''TFViTMAEModel''',
		    '''TFViTMAEPreTrainedModel''',
		]

if TYPE_CHECKING:
		from .configuration_vit_mae import VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTMAEConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_vit_mae import (
				    VIT_MAE_PRETRAINED_MODEL_ARCHIVE_LIST,
				    ViTMAEForPreTraining,
				    ViTMAELayer,
				    ViTMAEModel,
				    ViTMAEPreTrainedModel,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_vit_mae import TFViTMAEForPreTraining, TFViTMAEModel, TFViTMAEPreTrainedModel


else:
		import sys

		__lowerCamelCase					:	Any       =  _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
