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




					def 				UpperCamelCase_    (			self			:       Dict    ):


							"""simple docstring"""




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

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




					def 				UpperCamelCase_    (			self			:       int    ):


							"""simple docstring"""




							__snake_case									=  FlaxStableDiffusionPipeline.from_pretrained(
							    "hf-internal-testing/tiny-stable-diffusion-pipe"   ,safety_checker=_lowerCAmelCase    )

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

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

							__snake_case									=  jax.device_count()
							__snake_case									=  num_samples * [prompt]
							__snake_case									=  pipeline.prepare_inputs(_lowerCAmelCase    )

							# shard inputs and rng
							__snake_case									=  replicate(_lowerCAmelCase    )
							__snake_case									=  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase    )
							__snake_case									=  shard(_lowerCAmelCase    )

							__snake_case									=  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_5_1_4_7_4_5    ) < 1E-3
									assert np.abs(np.abs(_lowerCAmelCase   ,dtype=np.floataa    ).sum() - 49_947.875    ) < 5E-1

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




					def 				UpperCamelCase_    (			self			:       Tuple    ):


							"""simple docstring"""




							__snake_case									=  FlaxStableDiffusionPipeline.from_pretrained(
							    "CompVis/stable-diffusion-v1-4"   ,revision="flax"   ,safety_checker=_lowerCAmelCase    )

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

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

							__snake_case									=  jax.device_count()
							__snake_case									=  num_samples * [prompt]
							__snake_case									=  pipeline.prepare_inputs(_lowerCAmelCase    )

							# shard inputs and rng
							__snake_case									=  replicate(_lowerCAmelCase    )
							__snake_case									=  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase    )
							__snake_case									=  shard(_lowerCAmelCase    )

							__snake_case									=  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.0_5_6_5_2_4_0_1)    ) < 1E-3
									assert np.abs((np.abs(_lowerCAmelCase   ,dtype=np.floataa    ).sum() - 2_383_808.2)    ) < 5E-1




					def 				UpperCamelCase_    (			self			:       List[str]    ):


							"""simple docstring"""




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

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

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

							__snake_case									=  jax.device_count()
							__snake_case									=  num_samples * [prompt]
							__snake_case									=  pipeline.prepare_inputs(_lowerCAmelCase    )

							# shard inputs and rng
							__snake_case									=  replicate(_lowerCAmelCase    )
							__snake_case									=  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase    )
							__snake_case									=  shard(_lowerCAmelCase    )

							__snake_case									=  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.0_4_0_0_3_9_0_6)    ) < 1E-3
									assert np.abs((np.abs(_lowerCAmelCase   ,dtype=np.floataa    ).sum() - 2_373_516.75)    ) < 5E-1




					def 				UpperCamelCase_    (			self			:       Dict    ):


							"""simple docstring"""




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

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

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

							__snake_case									=  jax.device_count()
							__snake_case									=  num_samples * [prompt]
							__snake_case									=  pipeline.prepare_inputs(_lowerCAmelCase    )

							# shard inputs and rng
							__snake_case									=  replicate(_lowerCAmelCase    )
							__snake_case									=  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase    )
							__snake_case									=  shard(_lowerCAmelCase    )

							__snake_case									=  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.0_4_0_0_3_9_0_6)    ) < 1E-3
									assert np.abs((np.abs(_lowerCAmelCase   ,dtype=np.floataa    ).sum() - 2_373_516.75)    ) < 5E-1




					def 				UpperCamelCase_    (			self			:       Dict    ):


							"""simple docstring"""




							__snake_case									=  FlaxDDIMScheduler(
							    beta_start=0.0_0_0_8_5   ,beta_end=0.0_1_2   ,beta_schedule="scaled_linear"   ,set_alpha_to_one=_lowerCAmelCase   ,steps_offset=1   ,)

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

							__snake_case									=  scheduler_state

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

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

							__snake_case									=  jax.device_count()
							__snake_case									=  num_samples * [prompt]
							__snake_case									=  pipeline.prepare_inputs(_lowerCAmelCase    )

							# shard inputs and rng
							__snake_case									=  replicate(_lowerCAmelCase    )
							__snake_case									=  jax.random.split(_lowerCAmelCase   ,_lowerCAmelCase    )
							__snake_case									=  shard(_lowerCAmelCase    )

							__snake_case									=  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.0_4_5_0_4_3_9_4_5)    ) < 1E-3
									assert np.abs((np.abs(_lowerCAmelCase   ,dtype=np.floataa    ).sum() - 2_347_693.5)    ) < 5E-1







					def 				UpperCamelCase_    (			self			:       Optional[int]    ):


							"""simple docstring"""




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

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

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

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

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

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







# Lint as: python3
import dataclasses
import re
from dataclasses import dataclass
from functools import total_ordering
from typing import Optional, Union


a_							=    re.compile(r"""^(?P<major>\d+)""" r"""\.(?P<minor>\d+)""" r"""\.(?P<patch>\d+)$""")







@total_ordering
@dataclass
class 					A_:







			"""simple docstring"""



			a_       :			str
			a_       :			Optional[str]												=			None
			a_       :			Optional[Union[str, int]]												=			None
			a_       :			Optional[Union[str, int]]												=			None
			a_       :			Optional[Union[str, int]]												=			None
			def    _lowerCAmelCase		(	self							):
						_lowerCamelCase       ,							_lowerCamelCase       ,							_lowerCamelCase					:					List[str]	       =    _str_to_version_tuple(self.version_str							)
			def __repr__(	self							):
						return F"{self.tuple[0]}.{self.tuple[1]}.{self.tuple[2]}"
			@property
			def    _lowerCAmelCase		(	self							):
						return self.major, self.minor, self.patch
			def    _lowerCAmelCase		(	self				,   A							):
						if isinstance(A				,   A							):
									return Version(A							)
						elif isinstance(A				,   A							):
									return other
						raise TypeError(F"{other} (type {type(A							)}) cannot be compared to version."							)
			def __eq__(	self				,   A							):
						try:
									_lowerCamelCase					:					Dict	       =    self._validate_operand(A							)
						except (TypeError, ValueError):
									return False
						else:
									return self.tuple == other.tuple
			def __lt__(	self				,   A							):
						_lowerCamelCase					:					Any	       =    self._validate_operand(A							)
						return self.tuple < other.tuple
			def __hash__(	self							):
						return hash(_version_tuple_to_str(self.tuple							)							)
			@classmethod
			def    _lowerCAmelCase		(	cls				,   A							):
						_lowerCamelCase					:					Optional[Any]	       =    {f.name for f in dataclasses.fields(cls							)}
						return cls(**{k: v for k, v in dic.items() if k in field_names}							)

			def    _lowerCAmelCase		(	self							):
						return self.version_str






def     UpperCAmelCase_					(				__a   :			Any						):




			'''simple docstring'''



			_lowerCamelCase					:					Union[str, Any]	       =    _VERSION_REG.match(__a						)
			if not res:
						raise ValueError(f"Invalid version '{version_str}'. Format should be x.y.z with {{x,y,z}} being digits."						)
			return tuple(int(__a						) for v in [res.group('major'						), res.group('minor'						), res.group('patch'						)]						)



def     UpperCAmelCase_					(				__a   :			List[str]						):




			'''simple docstring'''



			return ".".join(str(__a						) for v in version_tuple						)





