






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

					'''simple docstring'''

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


												"""simple docstring"""

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

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

					'''simple docstring'''

					def       lowerCAmelCase				(  self				:	List[Any]							)			->  Tuple:


												"""simple docstring"""

												__lowercase			,			__lowercase      :							Union[str, Any]				=     FlaxStableDiffusionPipeline.from_pretrained(
												    """hf-internal-testing/tiny-stable-diffusion-pipe"""		,     safety_checker=__a							)

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

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

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

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

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

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



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


												"""simple docstring"""

												__lowercase			,			__lowercase      :							Any				=     FlaxStableDiffusionPipeline.from_pretrained(
												    """CompVis/stable-diffusion-v1-4"""		,     revision="""flax"""		,     safety_checker=__a							)

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

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

												__lowercase      :							Optional[Any]				=     jax.device_count()
												__lowercase      :							Tuple				=     num_samples * [prompt]
												__lowercase      :							List[str]				=     pipeline.prepare_inputs(__a							)

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

												__lowercase      :							Optional[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.05652401)							) < 1E-3
																			assert np.abs((np.abs(__a		,     dtype=np.floataa							).sum() - 2383808.2)							) < 5E-1



					def       lowerCAmelCase				(  self				:	List[Any]							)			->  int:


												"""simple docstring"""

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

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

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

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

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

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



					def       lowerCAmelCase				(  self				:	Dict							)			->  int:


												"""simple docstring"""

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

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

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

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

												# shard inputs and rng
												__lowercase      :							Dict				=     replicate(__a							)
												__lowercase      :							str				=     jax.random.split(__a		,     __a							)
												__lowercase      :							Dict				=     shard(__a							)

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



					def       lowerCAmelCase				(  self				:	Dict							)			->  Union[str, Any]:


												"""simple docstring"""

												__lowercase      :							List[Any]				=     FlaxDDIMScheduler(
												    beta_start=0.00085		,     beta_end=0.012		,     beta_schedule="""scaled_linear"""		,     set_alpha_to_one=__a		,     steps_offset=1		,     )

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

												__lowercase      :							List[Any]				=     scheduler_state

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

												__lowercase      :							Optional[Any]				=     jax.random.PRNGKey(0							)
												__lowercase      :							Dict				=     50

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

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

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



					def       lowerCAmelCase				(  self				:	Dict							)			->  Tuple:


												"""simple docstring"""

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

												__lowercase      :							Optional[Any]				=     jax.device_count()
												__lowercase      :							Optional[int]				=     num_samples * [prompt]
												__lowercase      :							Union[str, Any]				=     jax.random.split(jax.random.PRNGKey(0							)		,     __a							)

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

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

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

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







import logging
import os
import threading
import time


try:
	import warnings
except ImportError:
	lowerCamelCase : Any						   =       None

try:
	import msvcrt
except ImportError:
	lowerCamelCase : str						   =       None

try:
	import fcntl
except ImportError:
	lowerCamelCase : Optional[Any]						   =       None


# Backward compatibility
# ------------------------------------------------
try:
	TimeoutError
except NameError:
	lowerCamelCase : Union[str, Any]						   =       OSError


# Data
# ------------------------------------------------
lowerCamelCase : Tuple						   =       [
    '''Timeout''',
    '''BaseFileLock''',
    '''WindowsFileLock''',
    '''UnixFileLock''',
    '''SoftFileLock''',
    '''FileLock''',
]

lowerCamelCase : Tuple						   =       '''3.0.12'''


lowerCamelCase : Any						   =       None
def       snake_case_		(     ):
							global _logger
							__lowercase      :							List[str]				=     _logger or logging.getLogger(__name__      )
							return _logger

class 		lowerCAmelCase     ( __a			):

					'''simple docstring'''

					def __init__(  self				:	Any		,     __a				:	Any							)			->  List[Any]:


												"""simple docstring"""

												__lowercase      :							List[str]				=     lock_file
												return None



					def __str__(  self				:	str							)			->  Any:


												"""simple docstring"""

												__lowercase      :							Any				=     F"The file lock '{self.lock_file}' could not be acquired."
												return temp







class 		lowerCAmelCase     :

					'''simple docstring'''

					def __init__(  self				:	List[Any]		,     __a				:	Optional[int]							)			->  int:


												"""simple docstring"""

												__lowercase      :							Optional[Any]				=     lock
												return None



					def __enter__(  self				:	Dict							)			->  Dict:


												"""simple docstring"""

												return self.lock



					def __exit__(  self				:	Optional[int]		,     __a				:	Dict		,     __a				:	Any		,     __a				:	Tuple							)			->  Optional[Any]:


												"""simple docstring"""

												self.lock.release()
												return None







class 		lowerCAmelCase     :

					'''simple docstring'''

					def __init__(  self				:	Tuple		,     __a				:	Any		,     __a				:	Dict=-1		,     __a				:	Optional[Any]=None							)			->  Any:


												"""simple docstring"""

												__lowercase      :							Optional[int]				=     max_filename_length if max_filename_length is not None else 255
												# Hash the filename if it's too long
												__lowercase      :							Dict				=     self.hash_filename_if_too_long(__a		,     __a							)
												# The path to the lock file.
												__lowercase      :							Optional[Any]				=     lock_file

												# The file descriptor for the *_lock_file* as it is returned by the
												# os.open() function.
												# This file lock is only NOT None, if the object currently holds the
												# lock.
												__lowercase      :							int				=     None

												# The default timeout value.
												__lowercase      :							Optional[int]				=     timeout

												# We use this lock primarily for the lock counter.
												__lowercase      :							Optional[Any]				=     threading.Lock()

												# The lock counter is used for implementing the nested locking
												# mechanism. Whenever the lock is acquired, the counter is increased and
												# the lock is only released, when this value is 0 again.
												__lowercase      :							Union[str, Any]				=     0
												return None



					@property
					def       lowerCAmelCase				(  self				:	Optional[Any]							)			->  Optional[int]:


												"""simple docstring"""

												return self._lock_file



					@property
					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  Optional[Any]:


												"""simple docstring"""

												return self._timeout



					@timeout.setter
					def       lowerCAmelCase				(  self				:	Tuple		,     __a				:	Tuple							)			->  Dict:


												"""simple docstring"""

												__lowercase      :							Tuple				=     float(__a							)
												return None



					def       lowerCAmelCase				(  self				:	Dict							)			->  Tuple:


												"""simple docstring"""

												raise NotImplementedError()



					def       lowerCAmelCase				(  self				:	Dict							)			->  Any:


												"""simple docstring"""

												raise NotImplementedError()



					@property
					def       lowerCAmelCase				(  self				:	Dict							)			->  List[Any]:


												"""simple docstring"""

												return self._lock_file_fd is not None



					def       lowerCAmelCase				(  self				:	Any		,     __a				:	Optional[Any]=None		,     __a				:	Union[str, Any]=0.05							)			->  List[str]:


												"""simple docstring"""

												if timeout is None:
																			__lowercase      :							Union[str, Any]				=     self.timeout

												# Increment the number right at the beginning.
												# We can still undo it, if something fails.
												with self._thread_lock:
																			self._lock_counter += 1

												__lowercase      :							int				=     id(self							)
												__lowercase      :							Optional[Any]				=     self._lock_file
												__lowercase      :							List[str]				=     time.time()
												try:
																			while True:
																										with self._thread_lock:
																																	if not self.is_locked:
																																								logger().debug(F"Attempting to acquire lock {lock_id} on {lock_filename}"							)
																																								self._acquire()

																										if self.is_locked:
																																	logger().debug(F"Lock {lock_id} acquired on {lock_filename}"							)
																																	break
																										elif timeout >= 0 and time.time() - start_time > timeout:
																																	logger().debug(F"Timeout on acquiring lock {lock_id} on {lock_filename}"							)
																																	raise Timeout(self._lock_file							)
																										else:
																																	logger().debug(
																																	    F"Lock {lock_id} not acquired on {lock_filename}, waiting {poll_intervall} seconds ..."							)
																																	time.sleep(__a							)
												except:  # noqa
																			# Something did go wrong, so decrement the counter.
																			with self._thread_lock:
																										__lowercase      :							Optional[int]				=     max(0		,     self._lock_counter - 1							)

																			raise
												return _Acquire_ReturnProxy(lock=self							)



					def       lowerCAmelCase				(  self				:	Union[str, Any]		,     __a				:	Optional[Any]=False							)			->  Optional[Any]:


												"""simple docstring"""

												with self._thread_lock:
																			if self.is_locked:
																										self._lock_counter -= 1

																										if self._lock_counter == 0 or force:
																																	__lowercase      :							Optional[Any]				=     id(self							)
																																	__lowercase      :							str				=     self._lock_file

																																	logger().debug(F"Attempting to release lock {lock_id} on {lock_filename}"							)
																																	self._release()
																																	__lowercase      :							List[str]				=     0
																																	logger().debug(F"Lock {lock_id} released on {lock_filename}"							)

												return None



					def __enter__(  self				:	Any							)			->  Optional[Any]:


												"""simple docstring"""

												self.acquire()
												return self



					def __exit__(  self				:	List[str]		,     __a				:	str		,     __a				:	int		,     __a				:	List[Any]							)			->  Tuple:


												"""simple docstring"""

												self.release()
												return None



					def __del__(  self				:	Optional[int]							)			->  Optional[int]:


												"""simple docstring"""

												self.release(force=__a							)
												return None



					def       lowerCAmelCase				(  self				:	Tuple		,     __a				:	str		,     __a				:	int							)			->  str:


												"""simple docstring"""

												__lowercase      :							List[Any]				=     os.path.basename(__a							)
												if len(__a							) > max_length and max_length > 0:
																			__lowercase      :							int				=     os.path.dirname(__a							)
																			__lowercase      :							List[str]				=     str(hash(__a							)							)
																			__lowercase      :							Optional[Any]				=     filename[: max_length - len(__a							) - 8] + """...""" + hashed_filename + """.lock"""
																			return os.path.join(__a		,     __a							)
												else:
																			return path







class 		lowerCAmelCase     ( __a			):

					'''simple docstring'''

					def __init__(  self				:	Union[str, Any]		,     __a				:	List[Any]		,     __a				:	Optional[int]=-1		,     __a				:	Tuple=None							)			->  List[Any]:


												"""simple docstring"""

												from .file_utils import relative_to_absolute_path

												super().__init__(__a		,     timeout=__a		,     max_filename_length=__a							)
												__lowercase      :							Tuple				=     """\\\\?\\""" + relative_to_absolute_path(self.lock_file							)



					def       lowerCAmelCase				(  self				:	Dict							)			->  Union[str, Any]:


												"""simple docstring"""

												__lowercase      :							List[str]				=     os.O_RDWR | os.O_CREAT | os.O_TRUNC

												try:
																			__lowercase      :							Tuple				=     os.open(self._lock_file		,     __a							)
												except OSError:
																			pass
												else:
																			try:
																										msvcrt.locking(__a		,     msvcrt.LK_NBLCK		,     1							)
																			except OSError:
																										os.close(__a							)
																			else:
																										__lowercase      :							Union[str, Any]				=     fd
												return None



					def       lowerCAmelCase				(  self				:	Any							)			->  List[Any]:


												"""simple docstring"""

												__lowercase      :							Optional[Any]				=     self._lock_file_fd
												__lowercase      :							int				=     None
												msvcrt.locking(__a		,     msvcrt.LK_UNLCK		,     1							)
												os.close(__a							)

												try:
																			os.remove(self._lock_file							)
												# Probably another instance of the application
												# that acquired the file lock.
												except OSError:
																			pass
												return None







class 		lowerCAmelCase     ( __a			):

					'''simple docstring'''

					def __init__(  self				:	List[str]		,     __a				:	Optional[Any]		,     __a				:	str=-1		,     __a				:	List[str]=None							)			->  Any:


												"""simple docstring"""

												__lowercase      :							Dict				=     os.statvfs(os.path.dirname(__a							)							).f_namemax
												super().__init__(__a		,     timeout=__a		,     max_filename_length=__a							)



					def       lowerCAmelCase				(  self				:	Optional[Any]							)			->  Dict:


												"""simple docstring"""

												__lowercase      :							List[Any]				=     os.O_RDWR | os.O_CREAT | os.O_TRUNC
												__lowercase      :							List[str]				=     os.open(self._lock_file		,     __a							)

												try:
																			fcntl.flock(__a		,     fcntl.LOCK_EX | fcntl.LOCK_NB							)
												except OSError:
																			os.close(__a							)
												else:
																			__lowercase      :							str				=     fd
												return None



					def       lowerCAmelCase				(  self				:	List[Any]							)			->  Optional[int]:


												"""simple docstring"""

												__lowercase      :							Any				=     self._lock_file_fd
												__lowercase      :							List[str]				=     None
												fcntl.flock(__a		,     fcntl.LOCK_UN							)
												os.close(__a							)
												return None







class 		lowerCAmelCase     ( __a			):

					'''simple docstring'''

					def       lowerCAmelCase				(  self				:	List[str]							)			->  Union[str, Any]:


												"""simple docstring"""

												__lowercase      :							Tuple				=     os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_TRUNC
												try:
																			__lowercase      :							Union[str, Any]				=     os.open(self._lock_file		,     __a							)
												except OSError:
																			pass
												else:
																			__lowercase      :							Optional[int]				=     fd
												return None



					def       lowerCAmelCase				(  self				:	List[str]							)			->  Tuple:


												"""simple docstring"""

												os.close(self._lock_file_fd							)
												__lowercase      :							int				=     None

												try:
																			os.remove(self._lock_file							)
												# The file is already deleted and that's what we want.
												except OSError:
																			pass
												return None







lowerCamelCase : Optional[Any]						   =       None

if msvcrt:
	lowerCamelCase : List[Any]						   =       WindowsFileLock
elif fcntl:
	lowerCamelCase : List[Any]						   =       UnixFileLock
else:
	lowerCamelCase : Union[str, Any]						   =       SoftFileLock

	if warnings is not None:
		warnings.warn('''only soft file lock is available''')