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



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					with tempfile.TemporaryDirectory() as tmpdirname:
									# pipeline has Flax weights
									lowerCamelCase_						=			FlaxDiffusionPipeline.from_pretrained(
									    "hf-internal-testing/tiny-stable-diffusion-pipe"						,  safety_checker=lowercase						,  cache_dir=lowercase )

									lowerCamelCase_						=			[t[-1] for t in os.walk(os.path.join(lowercase						,  os.listdir(lowercase )[0]						,  "snapshots" ) )]
									lowerCamelCase_						=			[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						):



	def     SCREAMING_SNAKE_CASE_(    self )				->       Dict:
					lowerCamelCase_	,    lowerCamelCase_						=			FlaxStableDiffusionPipeline.from_pretrained(
					    "hf-internal-testing/tiny-stable-diffusion-pipe"						,  safety_checker=lowercase )

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

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

					lowerCamelCase_						=			jax.device_count()
					lowerCamelCase_						=			num_samples * [prompt]
					lowerCamelCase_						=			pipeline.prepare_inputs(lowercase )

					# shard inputs and rng
					lowerCamelCase_						=			replicate(lowercase )
					lowerCamelCase_						=			jax.random.split(lowercase						,  lowercase )
					lowerCamelCase_						=			shard(lowercase )

					lowerCamelCase_						=			pipeline(lowercase						,  lowercase						,  lowercase						,  lowercase						,  jit=lowercase ).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_5_1_4_7_4_5 ) < 1e-3
									assert np.abs(np.abs(lowercase						,  dtype=np.floataa ).sum() - 4_9_9_4_7.8_7_5 ) < 5e-1

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



	def     SCREAMING_SNAKE_CASE_(    self )				->       List[Any]:
					lowerCamelCase_	,    lowerCamelCase_						=			FlaxStableDiffusionPipeline.from_pretrained(
					    "CompVis/stable-diffusion-v1-4"						,  revision="flax"						,  safety_checker=lowercase )

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

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

					lowerCamelCase_						=			jax.device_count()
					lowerCamelCase_						=			num_samples * [prompt]
					lowerCamelCase_						=			pipeline.prepare_inputs(lowercase )

					# shard inputs and rng
					lowerCamelCase_						=			replicate(lowercase )
					lowerCamelCase_						=			jax.random.split(lowercase						,  lowercase )
					lowerCamelCase_						=			shard(lowercase )

					lowerCamelCase_						=			pipeline(lowercase						,  lowercase						,  lowercase						,  lowercase						,  jit=lowercase ).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_5_6_5_2_4_0_1) ) < 1e-3
									assert np.abs((np.abs(lowercase						,  dtype=np.floataa ).sum() - 2_3_8_3_8_0_8.2) ) < 5e-1



	def     SCREAMING_SNAKE_CASE_(    self )				->       str:
					lowerCamelCase_	,    lowerCamelCase_						=			FlaxStableDiffusionPipeline.from_pretrained(
					    "CompVis/stable-diffusion-v1-4"						,  revision="bf16"						,  dtype=jnp.bfloataa						,  safety_checker=lowercase )

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

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

					lowerCamelCase_						=			jax.device_count()
					lowerCamelCase_						=			num_samples * [prompt]
					lowerCamelCase_						=			pipeline.prepare_inputs(lowercase )

					# shard inputs and rng
					lowerCamelCase_						=			replicate(lowercase )
					lowerCamelCase_						=			jax.random.split(lowercase						,  lowercase )
					lowerCamelCase_						=			shard(lowercase )

					lowerCamelCase_						=			pipeline(lowercase						,  lowercase						,  lowercase						,  lowercase						,  jit=lowercase ).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_4_0_0_3_9_0_6) ) < 1e-3
									assert np.abs((np.abs(lowercase						,  dtype=np.floataa ).sum() - 2_3_7_3_5_1_6.7_5) ) < 5e-1



	def     SCREAMING_SNAKE_CASE_(    self )				->       int:
					lowerCamelCase_	,    lowerCamelCase_						=			FlaxStableDiffusionPipeline.from_pretrained(
					    "CompVis/stable-diffusion-v1-4"						,  revision="bf16"						,  dtype=jnp.bfloataa )

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

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

					lowerCamelCase_						=			jax.device_count()
					lowerCamelCase_						=			num_samples * [prompt]
					lowerCamelCase_						=			pipeline.prepare_inputs(lowercase )

					# shard inputs and rng
					lowerCamelCase_						=			replicate(lowercase )
					lowerCamelCase_						=			jax.random.split(lowercase						,  lowercase )
					lowerCamelCase_						=			shard(lowercase )

					lowerCamelCase_						=			pipeline(lowercase						,  lowercase						,  lowercase						,  lowercase						,  jit=lowercase ).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_4_0_0_3_9_0_6) ) < 1e-3
									assert np.abs((np.abs(lowercase						,  dtype=np.floataa ).sum() - 2_3_7_3_5_1_6.7_5) ) < 5e-1



	def     SCREAMING_SNAKE_CASE_(    self )				->       Dict:
					lowerCamelCase_						=			FlaxDDIMScheduler(
					    beta_start=0.0_0_0_8_5						,  beta_end=0.0_1_2						,  beta_schedule="scaled_linear"						,  set_alpha_to_one=lowercase						,  steps_offset=1						,  )

					lowerCamelCase_	,    lowerCamelCase_						=			FlaxStableDiffusionPipeline.from_pretrained(
					    "CompVis/stable-diffusion-v1-4"						,  revision="bf16"						,  dtype=jnp.bfloataa						,  scheduler=lowercase						,  safety_checker=lowercase						,  )
					lowerCamelCase_						=			scheduler.create_state()

					lowerCamelCase_						=			scheduler_state

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

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

					lowerCamelCase_						=			jax.device_count()
					lowerCamelCase_						=			num_samples * [prompt]
					lowerCamelCase_						=			pipeline.prepare_inputs(lowercase )

					# shard inputs and rng
					lowerCamelCase_						=			replicate(lowercase )
					lowerCamelCase_						=			jax.random.split(lowercase						,  lowercase )
					lowerCamelCase_						=			shard(lowercase )

					lowerCamelCase_						=			pipeline(lowercase						,  lowercase						,  lowercase						,  lowercase						,  jit=lowercase ).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_4_5_0_4_3_9_4_5) ) < 1e-3
									assert np.abs((np.abs(lowercase						,  dtype=np.floataa ).sum() - 2_3_4_7_6_9_3.5) ) < 5e-1



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

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

					lowerCamelCase_	,    lowerCamelCase_						=			FlaxStableDiffusionPipeline.from_pretrained(
					    "CompVis/stable-diffusion-v1-4"						,  revision="bf16"						,  dtype=jnp.bfloataa						,  safety_checker=lowercase						,  )

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

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

					lowerCamelCase_						=			replicate(lowercase )
					lowerCamelCase_						=			pipeline.prepare_inputs(lowercase )
					lowerCamelCase_						=			shard(lowercase )
					lowerCamelCase_						=			pipeline(lowercase						,  lowercase						,  lowercase						,  jit=lowercase ).images
					assert images_eff.shape == (num_samples, 1, 512, 512, 3)
					lowerCamelCase_						=			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'''
lowerCamelCase :      Any					 =     "\n# Transformers installation\n! pip install transformers datasets\n# To install from source instead of the last release, comment the command above and uncomment the following one.\n# ! pip install git+https://github.com/huggingface/transformers.git\n"

lowerCamelCase :      int					 =     [{"type": "code", "content": INSTALL_CONTENT}]
lowerCamelCase :      str					 =     {
    "{processor_class}": "FakeProcessorClass",
    "{model_class}": "FakeModelClass",
    "{object_class}": "FakeObjectClass",
}
