






'''simple docstring'''





import logging
import os
import threading
import time


try:
						import warnings
except ImportError:
						__UpperCamelCase      :      List[str]     						=						None

try:
						import msvcrt
except ImportError:
						__UpperCamelCase      :      Optional[Any]     						=						None

try:
						import fcntl
except ImportError:
						__UpperCamelCase      :      List[str]     						=						None


# Backward compatibility
# ------------------------------------------------
try:
						TimeoutError
except NameError:
						__UpperCamelCase      :      Tuple     						=						OSError


# Data
# ------------------------------------------------
__UpperCamelCase      :      Optional[int]     						=						[
    """Timeout""",
    """BaseFileLock""",
    """WindowsFileLock""",
    """UnixFileLock""",
    """SoftFileLock""",
    """FileLock""",
]

__UpperCamelCase      :      List[str]     						=						"""3.0.12"""


__UpperCamelCase      :      Tuple     						=						None





def 				snake_case  (     ):




	'''simple docstring'''







	global _logger
	__lowercase             =  _logger or logging.getLogger(__name__					)
	return _logger


class        __UpperCamelCase    (     _lowerCAmelCase   ):







	def __init__(  self      :							Dict       ,				_lowerCAmelCase      :							Union[str, Any]			)		-> Tuple:
		"""simple docstring"""







		__lowercase             =  lock_file
		return None







	def __str__(  self      :							Union[str, Any]			)		-> Optional[int]:
		"""simple docstring"""







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


class        __UpperCamelCase    :







	def __init__(  self      :							Optional[Any]       ,				_lowerCAmelCase      :							Union[str, Any]			)		-> Tuple:
		"""simple docstring"""







		__lowercase             =  lock
		return None







	def __enter__(  self      :							Optional[Any]			)		-> List[Any]:
		"""simple docstring"""







		return self.lock







	def __exit__(  self      :							Dict       ,				_lowerCAmelCase      :							List[Any]       ,				_lowerCAmelCase      :							List[str]       ,				_lowerCAmelCase      :							int			)		-> str:
		"""simple docstring"""







		self.lock.release()
		return None


class        __UpperCamelCase    :







	def __init__(  self      :							List[str]       ,				_lowerCAmelCase      :							List[str]       ,				_lowerCAmelCase      :							Optional[int]=-1       ,				_lowerCAmelCase      :							List[Any]=None			)		-> Any:
		"""simple docstring"""







		__lowercase             =  max_filename_length if max_filename_length is not None else 255
		# Hash the filename if it's too long
		__lowercase             =  self.hash_filename_if_too_long(UpperCamelCase__       ,				UpperCamelCase__			)
		# The path to the lock file.
		__lowercase             =  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             =  None

		# The default timeout value.
		__lowercase             =  timeout

		# We use this lock primarily for the lock counter.
		__lowercase             =  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             =  0
		return None







	@property
	def  _a   (  self      :							str			)		-> Tuple:
		"""simple docstring"""







		return self._lock_file







	@property
	def  _a   (  self      :							List[str]			)		-> List[Any]:
		"""simple docstring"""







		return self._timeout







	@timeout.setter
	def  _a   (  self      :							Tuple       ,				_lowerCAmelCase      :							Any			)		-> Any:
		"""simple docstring"""







		__lowercase             =  float(UpperCamelCase__			)
		return None







	def  _a   (  self      :							List[str]			)		-> Any:
		"""simple docstring"""







		raise NotImplementedError()







	def  _a   (  self      :							Tuple			)		-> List[Any]:
		"""simple docstring"""







		raise NotImplementedError()







	@property
	def  _a   (  self      :							Optional[Any]			)		-> Union[str, Any]:
		"""simple docstring"""







		return self._lock_file_fd is not None







	def  _a   (  self      :							List[Any]       ,				_lowerCAmelCase      :							List[str]=None       ,				_lowerCAmelCase      :							List[str]=0.05			)		-> List[str]:
		"""simple docstring"""







		if timeout is None:
			__lowercase             =  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             =  id(self			)
		__lowercase             =  self._lock_file
		__lowercase             =  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(UpperCamelCase__			)
		except:  # noqa
			# Something did go wrong, so decrement the counter.
			with self._thread_lock:
				__lowercase             =  max(0       ,				self._lock_counter - 1			)

			raise
		return _Acquire_ReturnProxy(lock=self			)







	def  _a   (  self      :							List[str]       ,				_lowerCAmelCase      :							Optional[Any]=False			)		-> Union[str, Any]:
		"""simple docstring"""







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

				if self._lock_counter == 0 or force:
					__lowercase             =  id(self			)
					__lowercase             =  self._lock_file

					logger().debug(F'Attempting to release lock {lock_id} on {lock_filename}'			)
					self._release()
					__lowercase             =  0
					logger().debug(F'Lock {lock_id} released on {lock_filename}'			)

		return None







	def __enter__(  self      :							List[Any]			)		-> Optional[int]:
		"""simple docstring"""







		self.acquire()
		return self







	def __exit__(  self      :							Any       ,				_lowerCAmelCase      :							Optional[Any]       ,				_lowerCAmelCase      :							List[Any]       ,				_lowerCAmelCase      :							List[Any]			)		-> List[Any]:
		"""simple docstring"""







		self.release()
		return None







	def __del__(  self      :							Optional[Any]			)		-> Union[str, Any]:
		"""simple docstring"""







		self.release(force=UpperCamelCase__			)
		return None







	def  _a   (  self      :							Any       ,				_lowerCAmelCase      :							Optional[int]       ,				_lowerCAmelCase      :							Tuple			)		-> str:
		"""simple docstring"""







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


class        __UpperCamelCase    (     _lowerCAmelCase   ):







	def __init__(  self      :							Tuple       ,				_lowerCAmelCase      :							Any       ,				_lowerCAmelCase      :							str=-1       ,				_lowerCAmelCase      :							Dict=None			)		-> Union[str, Any]:
		"""simple docstring"""







		from .file_utils import relative_to_absolute_path

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







	def  _a   (  self      :							str			)		-> Optional[int]:
		"""simple docstring"""







		__lowercase             =  os.O_RDWR | os.O_CREAT | os.O_TRUNC

		try:
			__lowercase             =  os.open(self._lock_file       ,				UpperCamelCase__			)
		except OSError:
			pass
		else:
			try:
				msvcrt.locking(UpperCamelCase__       ,				msvcrt.LK_NBLCK       ,				1			)
			except OSError:
				os.close(UpperCamelCase__			)
			else:
				__lowercase             =  fd
		return None







	def  _a   (  self      :							str			)		-> int:
		"""simple docstring"""







		__lowercase             =  self._lock_file_fd
		__lowercase             =  None
		msvcrt.locking(UpperCamelCase__       ,				msvcrt.LK_UNLCK       ,				1			)
		os.close(UpperCamelCase__			)

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


class        __UpperCamelCase    (     _lowerCAmelCase   ):







	def __init__(  self      :							List[Any]       ,				_lowerCAmelCase      :							Dict       ,				_lowerCAmelCase      :							List[str]=-1       ,				_lowerCAmelCase      :							int=None			)		-> Any:
		"""simple docstring"""







		__lowercase             =  os.statvfs(os.path.dirname(UpperCamelCase__			)			).f_namemax
		super().__init__(UpperCamelCase__       ,				timeout=UpperCamelCase__       ,				max_filename_length=UpperCamelCase__			)







	def  _a   (  self      :							str			)		-> str:
		"""simple docstring"""







		__lowercase             =  os.O_RDWR | os.O_CREAT | os.O_TRUNC
		__lowercase             =  os.open(self._lock_file       ,				UpperCamelCase__			)

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







	def  _a   (  self      :							Any			)		-> Tuple:
		"""simple docstring"""







		__lowercase             =  self._lock_file_fd
		__lowercase             =  None
		fcntl.flock(UpperCamelCase__       ,				fcntl.LOCK_UN			)
		os.close(UpperCamelCase__			)
		return None


class        __UpperCamelCase    (     _lowerCAmelCase   ):







	def  _a   (  self      :							Optional[int]			)		-> Tuple:
		"""simple docstring"""







		__lowercase             =  os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_TRUNC
		try:
			__lowercase             =  os.open(self._lock_file       ,				UpperCamelCase__			)
		except OSError:
			pass
		else:
			__lowercase             =  fd
		return None







	def  _a   (  self      :							str			)		-> List[str]:
		"""simple docstring"""







		os.close(self._lock_file_fd			)
		__lowercase             =  None

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

__UpperCamelCase      :      str     						=						None

if msvcrt:
						__UpperCamelCase      :      int     						=						WindowsFileLock
elif fcntl:
						__UpperCamelCase      :      Tuple     						=						UnixFileLock
else:
						__UpperCamelCase      :      List[str]     						=						SoftFileLock

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









import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
						import jax
						import jax.numpy as jnp
						from flax.jax_utils import replicate
						from flax.training.common_utils import shard


@slow
@require_flax
class        __UpperCamelCase    (     unittest.TestCase   ):







	def  _a   (  self      :							Union[str, Any]			)		-> Any:
		"""simple docstring"""







		super().tearDown()
		gc.collect()







	def  _a   (  self      :							List[str]			)		-> int:
		"""simple docstring"""







		__lowercase       ,						__lowercase             =  FlaxStableDiffusionPipeline.from_pretrained(
		    """stabilityai/stable-diffusion-2"""       ,				revision="""bf16"""       ,				dtype=jnp.bfloataa       ,				)

		__lowercase             =  """A painting of a squirrel eating a burger"""
		__lowercase             =  jax.device_count()
		__lowercase             =  num_samples * [prompt]
		__lowercase             =  sd_pipe.prepare_inputs(_lowerCAmelCase			)

		__lowercase             =  replicate(_lowerCAmelCase			)
		__lowercase             =  shard(_lowerCAmelCase			)

		__lowercase             =  jax.random.PRNGKey(0			)
		__lowercase             =  jax.random.split(_lowerCAmelCase       ,				jax.device_count()			)

		__lowercase             =  sd_pipe(_lowerCAmelCase       ,				_lowerCAmelCase       ,				_lowerCAmelCase       ,				num_inference_steps=25       ,				jit=_lowerCAmelCase			)[0]
		assert images.shape == (jax.device_count(), 1, 768, 768, 3)

		__lowercase             =  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]			)
		__lowercase             =  images[0, 253:256, 253:256, -1]

		__lowercase             =  jnp.asarray(jax.device_get(image_slice.flatten()			)			)
		__lowercase             =  jnp.array([0.4_238, 0.4_414, 0.4_395, 0.4_453, 0.4_629, 0.4_590, 0.4_531, 0.45_508, 0.4_512]			)
		print(F'output_slice: {output_slice}'			)
		assert jnp.abs(output_slice - expected_slice			).max() < 1e-2







	def  _a   (  self      :							str			)		-> List[Any]:
		"""simple docstring"""







		__lowercase             =  """stabilityai/stable-diffusion-2"""
		__lowercase       ,						__lowercase             =  FlaxDPMSolverMultistepScheduler.from_pretrained(_lowerCAmelCase       ,				subfolder="""scheduler"""			)
		__lowercase       ,						__lowercase             =  FlaxStableDiffusionPipeline.from_pretrained(
		    _lowerCAmelCase       ,				scheduler=_lowerCAmelCase       ,				revision="""bf16"""       ,				dtype=jnp.bfloataa       ,				)
		__lowercase             =  scheduler_params

		__lowercase             =  """A painting of a squirrel eating a burger"""
		__lowercase             =  jax.device_count()
		__lowercase             =  num_samples * [prompt]
		__lowercase             =  sd_pipe.prepare_inputs(_lowerCAmelCase			)

		__lowercase             =  replicate(_lowerCAmelCase			)
		__lowercase             =  shard(_lowerCAmelCase			)

		__lowercase             =  jax.random.PRNGKey(0			)
		__lowercase             =  jax.random.split(_lowerCAmelCase       ,				jax.device_count()			)

		__lowercase             =  sd_pipe(_lowerCAmelCase       ,				_lowerCAmelCase       ,				_lowerCAmelCase       ,				num_inference_steps=25       ,				jit=_lowerCAmelCase			)[0]
		assert images.shape == (jax.device_count(), 1, 768, 768, 3)

		__lowercase             =  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]			)
		__lowercase             =  images[0, 253:256, 253:256, -1]

		__lowercase             =  jnp.asarray(jax.device_get(image_slice.flatten()			)			)
		__lowercase             =  jnp.array([0.4_336, 0.42_969, 0.4_453, 0.4_199, 0.4_297, 0.4_531, 0.4_434, 0.4_434, 0.4_297]			)
		print(F'output_slice: {output_slice}'			)
		assert jnp.abs(output_slice - expected_slice			).max() < 1e-2

