


import logging
import os
import threading
import time


try:
			import warnings
except ImportError:
			UpperCAmelCase      =   None

try:
			import msvcrt
except ImportError:
			UpperCAmelCase      =   None

try:
			import fcntl
except ImportError:
			UpperCAmelCase      =   None


# Backward compatibility
# ------------------------------------------------
try:
			TimeoutError
except NameError:
			UpperCAmelCase      =   OSError


# Data
# ------------------------------------------------
UpperCAmelCase      =   [
    """Timeout""",
    """BaseFileLock""",
    """WindowsFileLock""",
    """UnixFileLock""",
    """SoftFileLock""",
    """FileLock""",
]

UpperCAmelCase      =   """3.0.12"""


UpperCAmelCase      =   None






def 	__lowerCAmelCase   ()->  List[Any]:




						"""simple docstring"""

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







class lowerCAmelCase_						(			lowerCamelCase__	):

	'''simple docstring'''

	def __init__(     self		,				_UpperCAmelCase ):
							snake_case_	=				lock_file
							return None





	def __str__(     self ):
							snake_case_	=				F'''The file lock \'{self.lock_file}\' could not be acquired.'''
							return temp







class lowerCAmelCase_						:

	'''simple docstring'''

	def __init__(     self		,				_UpperCAmelCase ):
							snake_case_	=				lock
							return None



	def __enter__(     self ):
							return self.lock





	def __exit__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
							self.lock.release()
							return None







class lowerCAmelCase_						:

	'''simple docstring'''

	def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=-1		,				_UpperCAmelCase=None ):
							snake_case_	=				max_filename_length if max_filename_length is not None else 2_55
							# Hash the filename if it's too long
							snake_case_	=				self.hash_filename_if_too_long(_UpperCAmelCase		,				_UpperCAmelCase )
							# The path to the lock file.
							snake_case_	=				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.
							snake_case_	=				None

							# The default timeout value.
							snake_case_	=				timeout

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



	@property
	def 						UpperCamelCase__					(     self ):
							return self._lock_file



	@property
	def 						UpperCamelCase__					(     self ):
							return self._timeout



	@timeout.setter
	def 						UpperCamelCase__					(     self		,				_UpperCAmelCase ):
							snake_case_	=				float(_UpperCAmelCase )
							return None



	def 						UpperCamelCase__					(     self ):
							raise NotImplementedError()



	def 						UpperCamelCase__					(     self ):
							raise NotImplementedError()



	@property
	def 						UpperCamelCase__					(     self ):
							return self._lock_file_fd is not None



	def 						UpperCamelCase__					(     self		,				_UpperCAmelCase=None		,				_UpperCAmelCase=0.05 ):
							# Use the default timeout, if no timeout is provided.
							if timeout is None:
													snake_case_	=				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

							snake_case_	=				id(self )
							snake_case_	=				self._lock_file
							snake_case_	=				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:
																			snake_case_	=				max(0		,				self._lock_counter - 1 )

													raise
							return _Acquire_ReturnProxy(lock=self )



	def 						UpperCamelCase__					(     self		,				_UpperCAmelCase=False ):
							with self._thread_lock:
													if self.is_locked:
																			self._lock_counter -= 1

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

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

							return None



	def __enter__(     self ):
							self.acquire()
							return self



	def __exit__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
							self.release()
							return None



	def __del__(     self ):
							self.release(force=_UpperCAmelCase )
							return None





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







class lowerCAmelCase_						(			lowerCamelCase__	):

	'''simple docstring'''

	def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=-1		,				_UpperCAmelCase=None ):
							from .file_utils import relative_to_absolute_path

							super().__init__(_UpperCAmelCase		,				timeout=_UpperCAmelCase		,				max_filename_length=_UpperCAmelCase )
							snake_case_	=				'''\\\\?\\''' + relative_to_absolute_path(self.lock_file )



	def 						UpperCamelCase__					(     self ):
							snake_case_	=				os.O_RDWR | os.O_CREAT | os.O_TRUNC

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





	def 						UpperCamelCase__					(     self ):
							snake_case_	=				self._lock_file_fd
							snake_case_	=				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 lowerCAmelCase_						(			lowerCamelCase__	):

	'''simple docstring'''

	def __init__(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=-1		,				_UpperCAmelCase=None ):
							snake_case_	=				os.statvfs(os.path.dirname(_UpperCAmelCase ) ).f_namemax
							super().__init__(_UpperCAmelCase		,				timeout=_UpperCAmelCase		,				max_filename_length=_UpperCAmelCase )



	def 						UpperCamelCase__					(     self ):
							snake_case_	=				os.O_RDWR | os.O_CREAT | os.O_TRUNC
							snake_case_	=				os.open(self._lock_file		,				_UpperCAmelCase )

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





	def 						UpperCamelCase__					(     self ):
							# Do not remove the lockfile:
							#
							#   https://github.com/benediktschmitt/py-filelock/issues/31
							#   https://stackoverflow.com/questions/17708885/flock-removing-locked-file-without-race-condition
							snake_case_	=				self._lock_file_fd
							snake_case_	=				None
							fcntl.flock(_UpperCAmelCase		,				fcntl.LOCK_UN )
							os.close(_UpperCAmelCase )
							return None







class lowerCAmelCase_						(			lowerCamelCase__	):

	'''simple docstring'''

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





	def 						UpperCamelCase__					(     self ):
							os.close(self._lock_file_fd )
							snake_case_	=				None

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





UpperCAmelCase      =   None

if msvcrt:
			UpperCAmelCase      =   WindowsFileLock
elif fcntl:
			UpperCAmelCase      =   UnixFileLock
else:
			UpperCAmelCase      =   SoftFileLock

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



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


UpperCAmelCase      =   False







class lowerCAmelCase_						(			unittest.TestCase	):

 '''simple docstring'''

 def 						UpperCamelCase__					(     self ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



 @property
 def 						UpperCamelCase__					(     self ):
       return 12



 @property
 def 						UpperCamelCase__					(     self ):
       return 12



 @property
 def 						UpperCamelCase__					(     self ):
       return 32



 @property
 def 						UpperCamelCase__					(     self ):
       torch.manual_seed(0 )
       snake_case_	=				VQModel(
           block_out_channels=[32, 64]		,				in_channels=3		,				out_channels=3		,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']		,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']		,				latent_channels=3		,				num_vq_embeddings=self.num_embed		,				vq_embed_dim=3		,				)
       return model



 @property
 def 						UpperCamelCase__					(     self ):
       snake_case_	=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )
       return tokenizer



 @property
 def 						UpperCamelCase__					(     self ):
       torch.manual_seed(0 )
       snake_case_	=				CLIPTextConfig(
           bos_token_id=0		,				eos_token_id=2		,				hidden_size=self.text_embedder_hidden_size		,				intermediate_size=37		,				layer_norm_eps=1E-05		,				num_attention_heads=4		,				num_hidden_layers=5		,				pad_token_id=1		,				vocab_size=10_00		,				)
       return CLIPTextModel(_UpperCAmelCase )



 @property
 def 						UpperCamelCase__					(     self ):
       torch.manual_seed(0 )

       snake_case_	=				12
       snake_case_	=				12

       snake_case_	=				{
           '''attention_bias''': True,
           '''cross_attention_dim''': 32,
           '''attention_head_dim''': height * width,
           '''num_attention_heads''': 1,
           '''num_vector_embeds''': self.num_embed,
           '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
           '''norm_num_groups''': 32,
           '''sample_size''': width,
           '''activation_fn''': '''geglu-approximate''',
       }

       snake_case_	=				TransformeraDModel(**_UpperCAmelCase )
       return model



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				'''cpu'''

       snake_case_	=				self.dummy_vqvae
       snake_case_	=				self.dummy_text_encoder
       snake_case_	=				self.dummy_tokenizer
       snake_case_	=				self.dummy_transformer
       snake_case_	=				VQDiffusionScheduler(self.num_embed )
       snake_case_	=				LearnedClassifierFreeSamplingEmbeddings(learnable=_UpperCAmelCase )

       snake_case_	=				VQDiffusionPipeline(
           vqvae=_UpperCAmelCase		,				text_encoder=_UpperCAmelCase		,				tokenizer=_UpperCAmelCase		,				transformer=_UpperCAmelCase		,				scheduler=_UpperCAmelCase		,				learned_classifier_free_sampling_embeddings=_UpperCAmelCase		,				)
       snake_case_	=				pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       snake_case_	=				'''teddy bear playing in the pool'''

       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipe([prompt]		,				generator=_UpperCAmelCase		,				num_inference_steps=2		,				output_type='''np''' )
       snake_case_	=				output.images

       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipe(
           [prompt]		,				generator=_UpperCAmelCase		,				output_type='''np'''		,				return_dict=_UpperCAmelCase		,				num_inference_steps=2 )[0]

       snake_case_	=				image[0, -3:, -3:, -1]
       snake_case_	=				image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       snake_case_	=				np.array([0.6_551, 0.6_168, 0.5_008, 0.5_676, 0.5_659, 0.4_295, 0.6_073, 0.5_599, 0.4_992] )

       assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2





 def 						UpperCamelCase__					(     self ):
       snake_case_	=				'''cpu'''

       snake_case_	=				self.dummy_vqvae
       snake_case_	=				self.dummy_text_encoder
       snake_case_	=				self.dummy_tokenizer
       snake_case_	=				self.dummy_transformer
       snake_case_	=				VQDiffusionScheduler(self.num_embed )
       snake_case_	=				LearnedClassifierFreeSamplingEmbeddings(
           learnable=_UpperCAmelCase		,				hidden_size=self.text_embedder_hidden_size		,				length=tokenizer.model_max_length )

       snake_case_	=				VQDiffusionPipeline(
           vqvae=_UpperCAmelCase		,				text_encoder=_UpperCAmelCase		,				tokenizer=_UpperCAmelCase		,				transformer=_UpperCAmelCase		,				scheduler=_UpperCAmelCase		,				learned_classifier_free_sampling_embeddings=_UpperCAmelCase		,				)
       snake_case_	=				pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       snake_case_	=				'''teddy bear playing in the pool'''

       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipe([prompt]		,				generator=_UpperCAmelCase		,				num_inference_steps=2		,				output_type='''np''' )
       snake_case_	=				output.images

       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipe(
           [prompt]		,				generator=_UpperCAmelCase		,				output_type='''np'''		,				return_dict=_UpperCAmelCase		,				num_inference_steps=2 )[0]

       snake_case_	=				image[0, -3:, -3:, -1]
       snake_case_	=				image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       snake_case_	=				np.array([0.6_693, 0.6_075, 0.4_959, 0.5_701, 0.5_583, 0.4_333, 0.6_171, 0.5_684, 0.4_988] )

       assert np.abs(image_slice.flatten() - expected_slice ).max() < 2.0
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2







@slow
@require_torch_gpu
class lowerCAmelCase_						(			unittest.TestCase	):

 '''simple docstring'''

 def 						UpperCamelCase__					(     self ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()





 def 						UpperCamelCase__					(     self ):
       snake_case_	=				load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy''' )

       snake_case_	=				VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq''' )
       snake_case_	=				pipeline.to(_UpperCAmelCase )
       pipeline.set_progress_bar_config(disable=_UpperCAmelCase )

       # requires GPU generator for gumbel softmax
       # don't use GPU generator in tests though
       snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
       snake_case_	=				pipeline(
           '''teddy bear playing in the pool'''		,				num_images_per_prompt=1		,				generator=_UpperCAmelCase		,				output_type='''np'''		,				)

       snake_case_	=				output.images[0]

       assert image.shape == (2_56, 2_56, 3)
       assert np.abs(expected_image - image ).max() < 2.0