





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    UpperCAmelCase			(						self							:     int   )       -> Optional[Any]:

			'''simple docstring'''




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

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

			'''simple docstring'''




			__lowercase						,      __lowercase											=     FlaxStableDiffusionPipeline.from_pretrained(
			    'hf-internal-testing/tiny-stable-diffusion-pipe'					,      safety_checker=__lowerCamelCase   )

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

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

			__lowercase											=     jax.device_count()
			__lowercase											=     num_samples * [prompt]
			__lowercase											=     pipeline.prepare_inputs(__lowerCamelCase   )

			# shard inputs and rng
			__lowercase											=     replicate(__lowerCamelCase   )
			__lowercase											=     jax.random.split(__lowerCamelCase					,      __lowerCamelCase   )
			__lowercase											=     shard(__lowerCamelCase   )

			__lowercase											=     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.1514745   ) < 1E-3
					assert np.abs(np.abs(__lowerCamelCase					,      dtype=np.floataa   ).sum() - 49947.875   ) < 5E-1

			__lowercase											=     pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]   )   )   )
			assert len(__lowerCamelCase   ) == num_samples




	def    UpperCAmelCase			(						self							:     Any   )       -> Any:

			'''simple docstring'''




			__lowercase						,      __lowercase											=     FlaxStableDiffusionPipeline.from_pretrained(
			    'CompVis/stable-diffusion-v1-4'					,      revision='flax'					,      safety_checker=__lowerCamelCase   )

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

			__lowercase											=     jax.random.PRNGKey(0   )
			__lowercase											=     50

			__lowercase											=     jax.device_count()
			__lowercase											=     num_samples * [prompt]
			__lowercase											=     pipeline.prepare_inputs(__lowerCamelCase   )

			# shard inputs and rng
			__lowercase											=     replicate(__lowerCamelCase   )
			__lowercase											=     jax.random.split(__lowerCamelCase					,      __lowerCamelCase   )
			__lowercase											=     shard(__lowerCamelCase   )

			__lowercase											=     pipeline(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      jit=__lowerCamelCase   ).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.05652401)   ) < 1E-3
					assert np.abs((np.abs(__lowerCamelCase					,      dtype=np.floataa   ).sum() - 2383808.2)   ) < 5E-1




	def    UpperCAmelCase			(						self							:     Tuple   )       -> Optional[Any]:

			'''simple docstring'''




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

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

			__lowercase											=     jax.random.PRNGKey(0   )
			__lowercase											=     50

			__lowercase											=     jax.device_count()
			__lowercase											=     num_samples * [prompt]
			__lowercase											=     pipeline.prepare_inputs(__lowerCamelCase   )

			# shard inputs and rng
			__lowercase											=     replicate(__lowerCamelCase   )
			__lowercase											=     jax.random.split(__lowerCamelCase					,      __lowerCamelCase   )
			__lowercase											=     shard(__lowerCamelCase   )

			__lowercase											=     pipeline(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      jit=__lowerCamelCase   ).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.04003906)   ) < 1E-3
					assert np.abs((np.abs(__lowerCamelCase					,      dtype=np.floataa   ).sum() - 2373516.75)   ) < 5E-1




	def    UpperCAmelCase			(						self							:     Dict   )       -> List[Any]:

			'''simple docstring'''




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

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

			__lowercase											=     jax.random.PRNGKey(0   )
			__lowercase											=     50

			__lowercase											=     jax.device_count()
			__lowercase											=     num_samples * [prompt]
			__lowercase											=     pipeline.prepare_inputs(__lowerCamelCase   )

			# shard inputs and rng
			__lowercase											=     replicate(__lowerCamelCase   )
			__lowercase											=     jax.random.split(__lowerCamelCase					,      __lowerCamelCase   )
			__lowercase											=     shard(__lowerCamelCase   )

			__lowercase											=     pipeline(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      jit=__lowerCamelCase   ).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.04003906)   ) < 1E-3
					assert np.abs((np.abs(__lowerCamelCase					,      dtype=np.floataa   ).sum() - 2373516.75)   ) < 5E-1




	def    UpperCAmelCase			(						self							:     List[str]   )       -> int:

			'''simple docstring'''




			__lowercase											=     FlaxDDIMScheduler(
			    beta_start=0.00085					,      beta_end=0.012					,      beta_schedule='scaled_linear'					,      set_alpha_to_one=__lowerCamelCase					,      steps_offset=1					,      )

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

			__lowercase											=     scheduler_state

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

			__lowercase											=     jax.random.PRNGKey(0   )
			__lowercase											=     50

			__lowercase											=     jax.device_count()
			__lowercase											=     num_samples * [prompt]
			__lowercase											=     pipeline.prepare_inputs(__lowerCamelCase   )

			# shard inputs and rng
			__lowercase											=     replicate(__lowerCamelCase   )
			__lowercase											=     jax.random.split(__lowerCamelCase					,      __lowerCamelCase   )
			__lowercase											=     shard(__lowerCamelCase   )

			__lowercase											=     pipeline(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      jit=__lowerCamelCase   ).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.045043945)   ) < 1E-3
					assert np.abs((np.abs(__lowerCamelCase					,      dtype=np.floataa   ).sum() - 2347693.5)   ) < 5E-1






	def    UpperCAmelCase			(						self							:     List[str]   )       -> List[Any]:

			'''simple docstring'''




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

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

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

			__lowercase											=     replicate(__lowerCamelCase   )
			__lowercase											=     pipeline.prepare_inputs(__lowerCamelCase   )
			__lowercase											=     shard(__lowerCamelCase   )
			__lowercase											=     pipeline(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      jit=__lowerCamelCase   ).images
			assert images.shape == (num_samples, 1, 512, 512, 3)
			__lowercase											=     images[2, 0, 256, 10:17, 1]

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

			__lowercase											=     replicate(__lowerCamelCase   )
			__lowercase											=     pipeline.prepare_inputs(__lowerCamelCase   )
			__lowercase											=     shard(__lowerCamelCase   )
			__lowercase											=     pipeline(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase					,      jit=__lowerCamelCase   ).images
			assert images_eff.shape == (num_samples, 1, 512, 512, 3)
			__lowercase											=     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











def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case = False							)		->    str:

		if not isinstance(snake_case						,   snake_case							):
				__lowercase											=     F"Expected string as input, found {type(snake_case							)}"
				raise ValueError(snake_case							)
		if not isinstance(snake_case						,   snake_case							):
				__lowercase											=     F"Expected boolean as use_pascal parameter, found {type(snake_case							)}"
				raise ValueError(snake_case							)

		__lowercase											=     input_str.split('_'							)

		__lowercase											=     0 if use_pascal else 1

		__lowercase											=     words[start_index:]

		__lowercase											=     [word[0].upper() + word[1:] for word in words_to_capitalize]

		__lowercase											=     '' if use_pascal else words[0]

		return "".join([initial_word, *capitalized_words]							)


if __name__ == "__main__":
		from doctest import testmod

		testmod()




