"""simple docstring"""
import gc
import unittest

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

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class 	lowerCAmelCase_     (      UpperCAmelCase_  ,     unittest.TestCase       ):



				"""simple docstring"""
				_lowerCAmelCase  :     Tuple      	=    UnCLIPPipeline
				_lowerCAmelCase  :     Any      	=    TEXT_TO_IMAGE_PARAMS - {
				    """negative_prompt""",
				    """height""",
				    """width""",
				    """negative_prompt_embeds""",
				    """guidance_scale""",
				    """prompt_embeds""",
				    """cross_attention_kwargs""",
				}
				_lowerCAmelCase  :     Tuple      	=    TEXT_TO_IMAGE_BATCH_PARAMS
				_lowerCAmelCase  :     List[str]      	=    [
				    """generator""",
				    """return_dict""",
				    """prior_num_inference_steps""",
				    """decoder_num_inference_steps""",
				    """super_res_num_inference_steps""",
				]
				_lowerCAmelCase  :     Union[str, Any]      	=    False

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										return 32

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										return 32

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										return self.time_input_dim

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										return self.time_input_dim * 4

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										return 1_00

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )
										return tokenizer

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										torch.manual_seed(0     )
										snake_case             =      CLIPTextConfig(
										    bos_token_id=0		,     eos_token_id=2		,     hidden_size=self.text_embedder_hidden_size		,     projection_dim=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 CLIPTextModelWithProjection(SCREAMING_SNAKE_CASE_     )

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										torch.manual_seed(0     )

										snake_case             =      {
										    '''num_attention_heads''': 2,
										    '''attention_head_dim''': 12,
										    '''embedding_dim''': self.text_embedder_hidden_size,
										    '''num_layers''': 1,
										}

										snake_case             =      PriorTransformer(**SCREAMING_SNAKE_CASE_     )
										return model

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										torch.manual_seed(0     )

										snake_case             =      {
										    '''clip_embeddings_dim''': self.text_embedder_hidden_size,
										    '''time_embed_dim''': self.time_embed_dim,
										    '''cross_attention_dim''': self.cross_attention_dim,
										}

										snake_case             =      UnCLIPTextProjModel(**SCREAMING_SNAKE_CASE_     )
										return model

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										torch.manual_seed(0     )

										snake_case             =      {
										    '''sample_size''': 32,
										    # RGB in channels
										    '''in_channels''': 3,
										    # Out channels is double in channels because predicts mean and variance
										    '''out_channels''': 6,
										    '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
										    '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
										    '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
										    '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
										    '''layers_per_block''': 1,
										    '''cross_attention_dim''': self.cross_attention_dim,
										    '''attention_head_dim''': 4,
										    '''resnet_time_scale_shift''': '''scale_shift''',
										    '''class_embed_type''': '''identity''',
										}

										snake_case             =      UNetaDConditionModel(**SCREAMING_SNAKE_CASE_     )
										return model

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										return {
										    "sample_size": 64,
										    "layers_per_block": 1,
										    "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
										    "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
										    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
										    "in_channels": 6,
										    "out_channels": 3,
										}

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										torch.manual_seed(0     )

										snake_case             =      UNetaDModel(**self.dummy_super_res_kwargs     )
										return model

				@property
				def 		snake_case			(					self     ):






										"""simple docstring"""

										torch.manual_seed(1     )

										snake_case             =      UNetaDModel(**self.dummy_super_res_kwargs     )
										return model

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      self.dummy_prior
										snake_case             =      self.dummy_decoder
										snake_case             =      self.dummy_text_proj
										snake_case             =      self.dummy_text_encoder
										snake_case             =      self.dummy_tokenizer
										snake_case             =      self.dummy_super_res_first
										snake_case             =      self.dummy_super_res_last

										snake_case             =      UnCLIPScheduler(
										    variance_type='fixed_small_log'		,     prediction_type='sample'		,     num_train_timesteps=10_00		,     clip_sample_range=5.0		,     )

										snake_case             =      UnCLIPScheduler(
										    variance_type='learned_range'		,     prediction_type='epsilon'		,     num_train_timesteps=10_00		,     )

										snake_case             =      UnCLIPScheduler(
										    variance_type='fixed_small_log'		,     prediction_type='epsilon'		,     num_train_timesteps=10_00		,     )

										snake_case             =      {
										    '''prior''': prior,
										    '''decoder''': decoder,
										    '''text_proj''': text_proj,
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''super_res_first''': super_res_first,
										    '''super_res_last''': super_res_last,
										    '''prior_scheduler''': prior_scheduler,
										    '''decoder_scheduler''': decoder_scheduler,
										    '''super_res_scheduler''': super_res_scheduler,
										}

										return components

				def 		snake_case			(					self		,     lowerCAmelCase		,     lowerCAmelCase=0     ):






										"""simple docstring"""

										if str(SCREAMING_SNAKE_CASE_     ).startswith('mps'     ):
																snake_case             =      torch.manual_seed(SCREAMING_SNAKE_CASE_     )
										else:
																snake_case             =      torch.Generator(device=SCREAMING_SNAKE_CASE_     ).manual_seed(SCREAMING_SNAKE_CASE_     )
										snake_case             =      {
										    '''prompt''': '''horse''',
										    '''generator''': generator,
										    '''prior_num_inference_steps''': 2,
										    '''decoder_num_inference_steps''': 2,
										    '''super_res_num_inference_steps''': 2,
										    '''output_type''': '''numpy''',
										}
										return inputs

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      '''cpu'''

										snake_case             =      self.get_dummy_components()

										snake_case             =      self.pipeline_class(**SCREAMING_SNAKE_CASE_     )
										snake_case             =      pipe.to(SCREAMING_SNAKE_CASE_     )

										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_     )

										snake_case             =      pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE_     )     )
										snake_case             =      output.images

										snake_case             =      pipe(
										    **self.get_dummy_inputs(SCREAMING_SNAKE_CASE_     )		,     return_dict=SCREAMING_SNAKE_CASE_		,     )[0]

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

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

										snake_case             =      np.array(
										    [
										        0.99_97,
										        0.99_88,
										        0.00_28,
										        0.99_97,
										        0.99_84,
										        0.99_65,
										        0.00_29,
										        0.99_86,
										        0.00_25,
										    ]     )

										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 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      torch.device('cpu'     )



										class 	lowerCAmelCase_     :



														"""simple docstring"""
														_lowerCAmelCase  :     List[Any]      	=    1





										snake_case             =      self.get_dummy_components()

										snake_case             =      self.pipeline_class(**SCREAMING_SNAKE_CASE_     )
										snake_case             =      pipe.to(SCREAMING_SNAKE_CASE_     )

										snake_case             =      components['''prior''']
										snake_case             =      components['''decoder''']
										snake_case             =      components['''super_res_first''']
										snake_case             =      components['''tokenizer''']
										snake_case             =      components['''text_encoder''']

										snake_case             =      torch.Generator(device=SCREAMING_SNAKE_CASE_     ).manual_seed(0     )
										snake_case             =      prior.dtype
										snake_case             =      1

										snake_case             =      (batch_size, prior.config.embedding_dim)
										snake_case             =      pipe.prepare_latents(
										    SCREAMING_SNAKE_CASE_		,     dtype=SCREAMING_SNAKE_CASE_		,     device=SCREAMING_SNAKE_CASE_		,     generator=SCREAMING_SNAKE_CASE_		,     latents=SCREAMING_SNAKE_CASE_		,     scheduler=DummyScheduler()     )
										snake_case             =      (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
										snake_case             =      pipe.prepare_latents(
										    SCREAMING_SNAKE_CASE_		,     dtype=SCREAMING_SNAKE_CASE_		,     device=SCREAMING_SNAKE_CASE_		,     generator=SCREAMING_SNAKE_CASE_		,     latents=SCREAMING_SNAKE_CASE_		,     scheduler=DummyScheduler()     )

										snake_case             =      (
										    batch_size,
										    super_res_first.config.in_channels // 2,
										    super_res_first.config.sample_size,
										    super_res_first.config.sample_size,
										)
										snake_case             =      pipe.prepare_latents(
										    SCREAMING_SNAKE_CASE_		,     dtype=SCREAMING_SNAKE_CASE_		,     device=SCREAMING_SNAKE_CASE_		,     generator=SCREAMING_SNAKE_CASE_		,     latents=SCREAMING_SNAKE_CASE_		,     scheduler=DummyScheduler()     )

										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_     )

										snake_case             =      '''this is a prompt example'''

										snake_case             =      torch.Generator(device=SCREAMING_SNAKE_CASE_     ).manual_seed(0     )
										snake_case             =      pipe(
										    [prompt]		,     generator=SCREAMING_SNAKE_CASE_		,     prior_num_inference_steps=2		,     decoder_num_inference_steps=2		,     super_res_num_inference_steps=2		,     prior_latents=SCREAMING_SNAKE_CASE_		,     decoder_latents=SCREAMING_SNAKE_CASE_		,     super_res_latents=SCREAMING_SNAKE_CASE_		,     output_type='np'		,     )
										snake_case             =      output.images

										snake_case             =      tokenizer(
										    SCREAMING_SNAKE_CASE_		,     padding='max_length'		,     max_length=tokenizer.model_max_length		,     return_tensors='pt'		,     )
										snake_case             =      text_encoder(text_inputs.input_ids     )
										snake_case             =      text_inputs.attention_mask

										snake_case             =      torch.Generator(device=SCREAMING_SNAKE_CASE_     ).manual_seed(0     )
										snake_case             =      pipe(
										    generator=SCREAMING_SNAKE_CASE_		,     prior_num_inference_steps=2		,     decoder_num_inference_steps=2		,     super_res_num_inference_steps=2		,     prior_latents=SCREAMING_SNAKE_CASE_		,     decoder_latents=SCREAMING_SNAKE_CASE_		,     super_res_latents=SCREAMING_SNAKE_CASE_		,     text_model_output=SCREAMING_SNAKE_CASE_		,     text_attention_mask=SCREAMING_SNAKE_CASE_		,     output_type='np'		,     )[0]

										# make sure passing text embeddings manually is identical
										assert np.abs(image - image_from_text     ).max() < 1E-4

				@skip_mps
				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      torch_device == '''cpu'''

										self._test_attention_slicing_forward_pass(test_max_difference=SCREAMING_SNAKE_CASE_		,     expected_max_diff=0.01     )

				@skip_mps
				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      torch_device == '''cpu'''
										snake_case             =      True
										snake_case             =      [
										    '''prior_num_inference_steps''',
										    '''decoder_num_inference_steps''',
										    '''super_res_num_inference_steps''',
										]

										self._test_inference_batch_single_identical(
										    test_max_difference=SCREAMING_SNAKE_CASE_		,     relax_max_difference=SCREAMING_SNAKE_CASE_		,     additional_params_copy_to_batched_inputs=SCREAMING_SNAKE_CASE_		,     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      [
										    '''prior_num_inference_steps''',
										    '''decoder_num_inference_steps''',
										    '''super_res_num_inference_steps''',
										]

										if torch_device == "mps":
																# TODO: MPS errors with larger batch sizes
																snake_case             =      [2, 3]
																self._test_inference_batch_consistent(
																    batch_sizes=SCREAMING_SNAKE_CASE_		,     additional_params_copy_to_batched_inputs=SCREAMING_SNAKE_CASE_		,     )
										else:
																self._test_inference_batch_consistent(
																    additional_params_copy_to_batched_inputs=SCREAMING_SNAKE_CASE_     )

				@skip_mps
				def 		snake_case			(					self     ):






										"""simple docstring"""

										return super().test_dict_tuple_outputs_equivalent()

				@skip_mps
				def 		snake_case			(					self     ):






										"""simple docstring"""

										return super().test_save_load_local()

				@skip_mps
				def 		snake_case			(					self     ):






										"""simple docstring"""

										return super().test_save_load_optional_components()





@nightly
class 	lowerCAmelCase_     (      unittest.TestCase       ):



				"""simple docstring"""

				def 		snake_case			(					self     ):






										"""simple docstring"""

										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      load_numpy(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										    '/unclip/karlo_v1_alpha_horse_cpu.npy'     )

										snake_case             =      UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'     )
										pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_     )

										snake_case             =      torch.manual_seed(0     )
										snake_case             =      pipeline(
										    'horse'		,     num_images_per_prompt=1		,     generator=SCREAMING_SNAKE_CASE_		,     output_type='np'		,     )

										snake_case             =      output.images[0]

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





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



				"""simple docstring"""

				def 		snake_case			(					self     ):






										"""simple docstring"""

										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()

				def 		snake_case			(					self     ):






										"""simple docstring"""

										snake_case             =      load_numpy(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										    '/unclip/karlo_v1_alpha_horse_fp16.npy'     )

										snake_case             =      UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'		,     torch_dtype=torch.floataa     )
										snake_case             =      pipeline.to(SCREAMING_SNAKE_CASE_     )
										pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_     )

										snake_case             =      torch.Generator(device='cpu'     ).manual_seed(0     )
										snake_case             =      pipeline(
										    'horse'		,     generator=SCREAMING_SNAKE_CASE_		,     output_type='np'		,     )

										snake_case             =      output.images[0]

										assert image.shape == (2_56, 2_56, 3)

										assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_		,     SCREAMING_SNAKE_CASE_     )

				def 		snake_case			(					self     ):






										"""simple docstring"""

										torch.cuda.empty_cache()
										torch.cuda.reset_max_memory_allocated()
										torch.cuda.reset_peak_memory_stats()

										snake_case             =      UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'		,     torch_dtype=torch.floataa     )
										snake_case             =      pipe.to(SCREAMING_SNAKE_CASE_     )
										pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_     )
										pipe.enable_attention_slicing()
										pipe.enable_sequential_cpu_offload()

										snake_case             =      pipe(
										    'horse'		,     num_images_per_prompt=1		,     prior_num_inference_steps=2		,     decoder_num_inference_steps=2		,     super_res_num_inference_steps=2		,     output_type='np'		,     )

										snake_case             =      torch.cuda.max_memory_allocated()
										# make sure that less than 7 GB is allocated
										assert mem_bytes < 7 * 10**9








'''simple docstring'''


import importlib
import importlib.metadata
import os
import warnings
from distutils.util import strtobool
from functools import lru_cache

import torch
from packaging import version
from packaging.version import parse

from .environment import parse_flag_from_env
from .versions import compare_versions, is_torch_version


try:
     import torch_xla.core.xla_model as xm  # noqa: F401

     UpperCamelCase			     =					True
except ImportError:
     UpperCamelCase			     =					False


# Cache this result has it's a C FFI call which can be pretty time-consuming
UpperCamelCase			     =					torch.distributed.is_available()
def 					SCREAMING_SNAKE_CASE(   __lowercase	)			-> Optional[int]:
     # Check we're not importing a "pkg_name" directory somewhere but the actual library by trying to grab the version
     A:		Optional[Any]          =		importlib.util.find_spec(__lowercase	) is not None
     if package_exists:
          try:
               A:		Dict          =		importlib.metadata.metadata(__lowercase	)
               return True
          except importlib.metadata.PackageNotFoundError:
               return False
def 					SCREAMING_SNAKE_CASE(   )			-> bool:
     return _torch_distributed_available
def 					SCREAMING_SNAKE_CASE(   )			-> Optional[int]:
     try:
          pass
     except ImportError:
          print(
              '''Intel(R) oneCCL Bindings for PyTorch* is required to run DDP on Intel(R) GPUs, but it is not'''
              ''' detected. If you see \"ValueError: Invalid backend: \'ccl\'\" error, please install Intel(R) oneCCL'''
              ''' Bindings for PyTorch*.'''	)
     return (
         importlib.util.find_spec('''torch_ccl'''	) is not None
         or importlib.util.find_spec('''oneccl_bindings_for_pytorch'''	) is not None
     )
def 					SCREAMING_SNAKE_CASE(   )			-> List[str]:
     return importlib.metadata.version('''oneccl_bind_pt'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Union[str, Any]:
     return _is_package_available('''transformer_engine'''	)
@lru_cache
def 					SCREAMING_SNAKE_CASE(   __lowercase=True	)			-> Union[str, Any]:
     # Due to bugs on the amp series GPUs, we disable torch-xla on them
     if torch.cuda.is_available():
          return False
     if _tpu_available and check_device:
          try:
               # Will raise a RuntimeError if no XLA configuration is found
               A:		Tuple          =		xm.xla_device()
               return True
          except RuntimeError:
               return False
     return _tpu_available
def 					SCREAMING_SNAKE_CASE(   )			-> List[str]:
     return _is_package_available('''deepspeed'''	)
def 					SCREAMING_SNAKE_CASE(   __lowercase=False	)			-> Any:
     if is_tpu_available():
          return not ignore_tpu
     if torch.cuda.is_available():
          return torch.cuda.is_bfaa_supported()
     if is_npu_available():
          return False
     return True
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     A:		List[str]          =		_is_package_available('''bitsandbytes'''	)
     if package_exists:
          A:		Union[str, Any]          =		version.parse(importlib.metadata.version('''bitsandbytes'''	)	)
          return compare_versions(__lowercase		,				'''>='''		,				'''0.39.0'''	)
     return False
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     A:		int          =		_is_package_available('''bitsandbytes'''	)
     if package_exists:
          A:		Tuple          =		version.parse(importlib.metadata.version('''bitsandbytes'''	)	)
          return compare_versions(__lowercase		,				'''>='''		,				'''0.37.2'''	)
     return False
def 					SCREAMING_SNAKE_CASE(   )			-> int:
     return _is_package_available('''bitsandbytes'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Optional[int]:
     if strtobool(os.environ.get('''ACCELERATE_USE_MEGATRON_LM'''		,				'''False'''	)	) == 1:
          A:		Any          =		importlib.util.find_spec('''megatron'''	) is not None
          if package_exists:
               try:
                    A:		str          =		parse(importlib.metadata.version('''megatron-lm'''	)	)
                    return compare_versions(__lowercase		,				'''>='''		,				'''2.2.0'''	)
               except Exception as e:
                    warnings.warn(F"""Parse Megatron version failed. Exception:{e}"""	)
                    return False
def 					SCREAMING_SNAKE_CASE(   )			-> Any:
     return _is_package_available('''safetensors'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Tuple:
     return _is_package_available('''transformers'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Union[str, Any]:
     return _is_package_available('''datasets'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Union[str, Any]:
     return _is_package_available('''aim'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> List[Any]:
     return _is_package_available('''tensorboard'''	) or _is_package_available('''tensorboardX'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Dict:
     return _is_package_available('''wandb'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Optional[Any]:
     return _is_package_available('''comet_ml'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Any:
     return _is_package_available('''boto3'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     if _is_package_available('''rich'''	):
          if "ACCELERATE_DISABLE_RICH" in os.environ:
               warnings.warn(
                   '''`ACCELERATE_DISABLE_RICH` is deprecated and will be removed in v0.22.0 and deactivated by default. Please use `ACCELERATE_ENABLE_RICH` if you wish to use `rich`.'''	)
               return not parse_flag_from_env('''ACCELERATE_DISABLE_RICH'''		,				__lowercase	)
          return parse_flag_from_env('''ACCELERATE_ENABLE_RICH'''		,				__lowercase	)
     return False
def 					SCREAMING_SNAKE_CASE(   )			-> int:
     return _is_package_available('''sagemaker'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> Union[str, Any]:
     return _is_package_available('''tqdm'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     return _is_package_available('''mlflow'''	)
def 					SCREAMING_SNAKE_CASE(   )			-> int:
     return is_torch_version('''>='''		,				'''1.12'''	) and torch.backends.mps.is_available() and torch.backends.mps.is_built()
def 					SCREAMING_SNAKE_CASE(   )			-> str:
     def get_major_and_minor_from_version(__lowercase	):
          return str(version.parse(__lowercase	).major	) + "." + str(version.parse(__lowercase	).minor	)

     A:		int          =		importlib.metadata.version('''torch'''	)
     if importlib.util.find_spec('''intel_extension_for_pytorch'''	) is None:
          return False
     A:		Optional[Any]          =		'''N/A'''
     try:
          A:		List[Any]          =		importlib.metadata.version('''intel_extension_for_pytorch'''	)
     except importlib.metadata.PackageNotFoundError:
          return False
     A:		Optional[int]          =		get_major_and_minor_from_version(_torch_version	)
     A:		List[str]          =		get_major_and_minor_from_version(_ipex_version	)
     if torch_major_and_minor != ipex_major_and_minor:
          warnings.warn(
              F"""Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*,"""
              F""" but PyTorch {_torch_version} is found. Please switch to the matching version and run again."""	)
          return False
     return True
@lru_cache
def 					SCREAMING_SNAKE_CASE(   __lowercase=False	)			-> Optional[Any]:
     if importlib.util.find_spec('''torch'''	) is None or importlib.util.find_spec('''torch_npu'''	) is None:
          return False

     import torch
     import torch_npu  # noqa: F401

     if check_device:
          try:
               # Will raise a RuntimeError if no NPU is found
               A:		Tuple          =		torch.npu.device_count()
               return torch.npu.is_available()
          except RuntimeError:
               return False
     return hasattr(__lowercase		,				'''npu'''	) and torch.npu.is_available()


@lru_cache
def 					SCREAMING_SNAKE_CASE(   __lowercase=False	)			-> Tuple:
     if not parse_flag_from_env('''ACCELERATE_USE_XPU'''		,				default=__lowercase	):
          return False
     if is_ipex_available():
          import torch

          if is_torch_version('''<='''		,				'''1.12'''	):
               return False
     else:
          return False

     import intel_extension_for_pytorch  # noqa: F401

     if check_device:
          try:
               # Will raise a RuntimeError if no XPU  is found
               A:		str          =		torch.xpu.device_count()
               return torch.xpu.is_available()
          except RuntimeError:
               return False
     return hasattr(__lowercase		,				'''xpu'''	) and torch.xpu.is_available()






