



'''simple docstring'''
import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetInpaintPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
    TEXT_TO_IMAGE_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()



class _A	(     snake_case_						,					snake_case_						,					snake_case_						,					unittest.TestCase     ):
		_SCREAMING_SNAKE_CASE       :     Any       =      StableDiffusionControlNetInpaintPipeline
		_SCREAMING_SNAKE_CASE       :     Dict       =      TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
		_SCREAMING_SNAKE_CASE       :     str       =      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
		_SCREAMING_SNAKE_CASE       :     Union[str, Any]       =      frozenset({"control_image"}     )  # skip `image` and `mask` for now, only test for control_image
		_SCREAMING_SNAKE_CASE       :     Tuple       =      TEXT_TO_IMAGE_IMAGE_PARAMS





		def 				__A							(  self					)							->  str:

					'''simple docstring'''
					torch.manual_seed(0					)
					__UpperCAmelCase :   Union[str, Any]    					=  UNetaDConditionModel(
					    block_out_channels=(32, 64)			,							layers_per_block=2			,							sample_size=32			,							in_channels=9			,							out_channels=4			,							down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,							up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")			,							cross_attention_dim=32			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   Any    					=  ControlNetModel(
					    block_out_channels=(32, 64)			,							layers_per_block=2			,							in_channels=4			,							down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,							cross_attention_dim=32			,							conditioning_embedding_out_channels=(16, 32)			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   Optional[int]    					=  DDIMScheduler(
					    beta_start=0.0_0085			,							beta_end=0.012			,							beta_schedule="""scaled_linear"""			,							clip_sample=lowercase_			,							set_alpha_to_one=lowercase_			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   Dict    					=  AutoencoderKL(
					    block_out_channels=[32, 64]			,							in_channels=3			,							out_channels=3			,							down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]			,							up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]			,							latent_channels=4			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   Optional[Any]    					=  CLIPTextConfig(
					    bos_token_id=0			,							eos_token_id=2			,							hidden_size=32			,							intermediate_size=37			,							layer_norm_eps=1E-05			,							num_attention_heads=4			,							num_hidden_layers=5			,							pad_token_id=1			,							vocab_size=1_000			,							)
					__UpperCAmelCase :   str    					=  CLIPTextModel(lowercase_					)
					__UpperCAmelCase :   Union[str, Any]    					=  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

					__UpperCAmelCase :   Dict    					=  {
					    'unet': unet,
					    'controlnet': controlnet,
					    'scheduler': scheduler,
					    'vae': vae,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					    'safety_checker': None,
					    'feature_extractor': None,
					}
					return components





		def 				__A							(  self			,							__UpperCAmelCase			,							__UpperCAmelCase=0					)							->  Optional[int]:

					'''simple docstring'''
					if str(lowercase_					).startswith("""mps"""					):
								__UpperCAmelCase :   Tuple    					=  torch.manual_seed(lowercase_					)
					else:
								__UpperCAmelCase :   str    					=  torch.Generator(device=lowercase_					).manual_seed(lowercase_					)

					__UpperCAmelCase :   str    					=  2
					__UpperCAmelCase :   str    					=  randn_tensor(
					    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)			,							generator=lowercase_			,							device=torch.device(lowercase_					)			,							)
					__UpperCAmelCase :   str    					=  floats_tensor((1, 3, 32, 32)			,							rng=random.Random(lowercase_					)					).to(lowercase_					)
					__UpperCAmelCase :   Tuple    					=  init_image.cpu().permute(0			,							2			,							3			,							1					)[0]

					__UpperCAmelCase :   int    					=  Image.fromarray(np.uinta(lowercase_					)					).convert("""RGB"""					).resize((64, 64)					)
					__UpperCAmelCase :   Union[str, Any]    					=  Image.fromarray(np.uinta(init_image + 4					)					).convert("""RGB"""					).resize((64, 64)					)

					__UpperCAmelCase :   int    					=  {
					    'prompt': 'A painting of a squirrel eating a burger',
					    'generator': generator,
					    'num_inference_steps': 2,
					    'guidance_scale': 6.0,
					    'output_type': 'numpy',
					    'image': image,
					    'mask_image': mask_image,
					    'control_image': control_image,
					}

					return inputs





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

					'''simple docstring'''
					return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3					)





		@unittest.skipIf(
		    torch_device != """cuda""" or not is_xformers_available()			,							reason="""XFormers attention is only available with CUDA and `xformers` installed"""			,							)
		def 				__A							(  self					)							->  List[str]:

					'''simple docstring'''
					self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3					)





		def 				__A							(  self					)							->  str:

					'''simple docstring'''
					self._test_inference_batch_single_identical(expected_max_diff=2E-3					)



class _A	(     snake_case_     ):
		_SCREAMING_SNAKE_CASE       :     Optional[int]       =      StableDiffusionControlNetInpaintPipeline
		_SCREAMING_SNAKE_CASE       :     str       =      TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
		_SCREAMING_SNAKE_CASE       :     Union[str, Any]       =      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
		_SCREAMING_SNAKE_CASE       :     Optional[int]       =      frozenset([]     )





		def 				__A							(  self					)							->  str:

					'''simple docstring'''
					torch.manual_seed(0					)
					__UpperCAmelCase :   Any    					=  UNetaDConditionModel(
					    block_out_channels=(32, 64)			,							layers_per_block=2			,							sample_size=32			,							in_channels=4			,							out_channels=4			,							down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,							up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")			,							cross_attention_dim=32			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   int    					=  ControlNetModel(
					    block_out_channels=(32, 64)			,							layers_per_block=2			,							in_channels=4			,							down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,							cross_attention_dim=32			,							conditioning_embedding_out_channels=(16, 32)			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   Optional[int]    					=  DDIMScheduler(
					    beta_start=0.0_0085			,							beta_end=0.012			,							beta_schedule="""scaled_linear"""			,							clip_sample=lowercase_			,							set_alpha_to_one=lowercase_			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   Dict    					=  AutoencoderKL(
					    block_out_channels=[32, 64]			,							in_channels=3			,							out_channels=3			,							down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]			,							up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]			,							latent_channels=4			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   Optional[Any]    					=  CLIPTextConfig(
					    bos_token_id=0			,							eos_token_id=2			,							hidden_size=32			,							intermediate_size=37			,							layer_norm_eps=1E-05			,							num_attention_heads=4			,							num_hidden_layers=5			,							pad_token_id=1			,							vocab_size=1_000			,							)
					__UpperCAmelCase :   Tuple    					=  CLIPTextModel(lowercase_					)
					__UpperCAmelCase :   str    					=  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

					__UpperCAmelCase :   int    					=  {
					    'unet': unet,
					    'controlnet': controlnet,
					    'scheduler': scheduler,
					    'vae': vae,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					    'safety_checker': None,
					    'feature_extractor': None,
					}
					return components



class _A	(     snake_case_						,					snake_case_						,					unittest.TestCase     ):
		_SCREAMING_SNAKE_CASE       :     Union[str, Any]       =      StableDiffusionControlNetInpaintPipeline
		_SCREAMING_SNAKE_CASE       :     Any       =      TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
		_SCREAMING_SNAKE_CASE       :     str       =      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS





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

					'''simple docstring'''
					torch.manual_seed(0					)
					__UpperCAmelCase :   List[Any]    					=  UNetaDConditionModel(
					    block_out_channels=(32, 64)			,							layers_per_block=2			,							sample_size=32			,							in_channels=9			,							out_channels=4			,							down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,							up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")			,							cross_attention_dim=32			,							)
					torch.manual_seed(0					)

					def init_weights(__UpperCAmelCase					):
								if isinstance(lowercase_			,							torch.nn.Convad					):
											torch.nn.init.normal(m.weight					)
											m.bias.data.fill_(1.0					)

					__UpperCAmelCase :   Any    					=  ControlNetModel(
					    block_out_channels=(32, 64)			,							layers_per_block=2			,							in_channels=4			,							down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,							cross_attention_dim=32			,							conditioning_embedding_out_channels=(16, 32)			,							)
					controlneta.controlnet_down_blocks.apply(lowercase_					)

					torch.manual_seed(0					)
					__UpperCAmelCase :   Any    					=  ControlNetModel(
					    block_out_channels=(32, 64)			,							layers_per_block=2			,							in_channels=4			,							down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,							cross_attention_dim=32			,							conditioning_embedding_out_channels=(16, 32)			,							)
					controlneta.controlnet_down_blocks.apply(lowercase_					)

					torch.manual_seed(0					)
					__UpperCAmelCase :   Union[str, Any]    					=  DDIMScheduler(
					    beta_start=0.0_0085			,							beta_end=0.012			,							beta_schedule="""scaled_linear"""			,							clip_sample=lowercase_			,							set_alpha_to_one=lowercase_			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   List[str]    					=  AutoencoderKL(
					    block_out_channels=[32, 64]			,							in_channels=3			,							out_channels=3			,							down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]			,							up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]			,							latent_channels=4			,							)
					torch.manual_seed(0					)
					__UpperCAmelCase :   List[str]    					=  CLIPTextConfig(
					    bos_token_id=0			,							eos_token_id=2			,							hidden_size=32			,							intermediate_size=37			,							layer_norm_eps=1E-05			,							num_attention_heads=4			,							num_hidden_layers=5			,							pad_token_id=1			,							vocab_size=1_000			,							)
					__UpperCAmelCase :   Optional[Any]    					=  CLIPTextModel(lowercase_					)
					__UpperCAmelCase :   Any    					=  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

					__UpperCAmelCase :   List[Any]    					=  MultiControlNetModel([controlneta, controlneta]					)

					__UpperCAmelCase :   List[Any]    					=  {
					    'unet': unet,
					    'controlnet': controlnet,
					    'scheduler': scheduler,
					    'vae': vae,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					    'safety_checker': None,
					    'feature_extractor': None,
					}
					return components





		def 				__A							(  self			,							__UpperCAmelCase			,							__UpperCAmelCase=0					)							->  List[Any]:

					'''simple docstring'''
					if str(lowercase_					).startswith("""mps"""					):
								__UpperCAmelCase :   Optional[int]    					=  torch.manual_seed(lowercase_					)
					else:
								__UpperCAmelCase :   Dict    					=  torch.Generator(device=lowercase_					).manual_seed(lowercase_					)

					__UpperCAmelCase :   Optional[int]    					=  2

					__UpperCAmelCase :   Dict    					=  [
					    randn_tensor(
					        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)			,							generator=lowercase_			,							device=torch.device(lowercase_					)			,							),
					    randn_tensor(
					        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)			,							generator=lowercase_			,							device=torch.device(lowercase_					)			,							),
					]
					__UpperCAmelCase :   Any    					=  floats_tensor((1, 3, 32, 32)			,							rng=random.Random(lowercase_					)					).to(lowercase_					)
					__UpperCAmelCase :   str    					=  init_image.cpu().permute(0			,							2			,							3			,							1					)[0]

					__UpperCAmelCase :   List[Any]    					=  Image.fromarray(np.uinta(lowercase_					)					).convert("""RGB"""					).resize((64, 64)					)
					__UpperCAmelCase :   Optional[int]    					=  Image.fromarray(np.uinta(init_image + 4					)					).convert("""RGB"""					).resize((64, 64)					)

					__UpperCAmelCase :   str    					=  {
					    'prompt': 'A painting of a squirrel eating a burger',
					    'generator': generator,
					    'num_inference_steps': 2,
					    'guidance_scale': 6.0,
					    'output_type': 'numpy',
					    'image': image,
					    'mask_image': mask_image,
					    'control_image': control_image,
					}

					return inputs





		def 				__A							(  self					)							->  int:

					'''simple docstring'''
					__UpperCAmelCase :   Tuple    					=  self.get_dummy_components()
					__UpperCAmelCase :   Optional[int]    					=  self.pipeline_class(**lowercase_					)
					pipe.to(lowercase_					)

					__UpperCAmelCase :   List[Any]    					=  10.0
					__UpperCAmelCase :   Optional[Any]    					=  4

					__UpperCAmelCase :   Dict    					=  self.get_dummy_inputs(lowercase_					)
					__UpperCAmelCase :   Optional[Any]    					=  steps
					__UpperCAmelCase :   Optional[Any]    					=  scale
					__UpperCAmelCase :   Dict    					=  pipe(**lowercase_					)[0]

					__UpperCAmelCase :   Optional[Any]    					=  self.get_dummy_inputs(lowercase_					)
					__UpperCAmelCase :   Any    					=  steps
					__UpperCAmelCase :   List[str]    					=  scale
					__UpperCAmelCase :   Tuple    					=  pipe(**lowercase_			,							control_guidance_start=0.1			,							control_guidance_end=0.2					)[0]

					__UpperCAmelCase :   List[Any]    					=  self.get_dummy_inputs(lowercase_					)
					__UpperCAmelCase :   Dict    					=  steps
					__UpperCAmelCase :   Any    					=  scale
					__UpperCAmelCase :   Tuple    					=  pipe(**lowercase_			,							control_guidance_start=[0.1, 0.3]			,							control_guidance_end=[0.2, 0.7]					)[0]

					__UpperCAmelCase :   int    					=  self.get_dummy_inputs(lowercase_					)
					__UpperCAmelCase :   List[str]    					=  steps
					__UpperCAmelCase :   Optional[Any]    					=  scale
					__UpperCAmelCase :   Dict    					=  pipe(**lowercase_			,							control_guidance_start=0.4			,							control_guidance_end=[0.5, 0.8]					)[0]

					# make sure that all outputs are different
					assert np.sum(np.abs(output_a - output_a					)					) > 1E-3
					assert np.sum(np.abs(output_a - output_a					)					) > 1E-3
					assert np.sum(np.abs(output_a - output_a					)					) > 1E-3





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

					'''simple docstring'''
					return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3					)





		@unittest.skipIf(
		    torch_device != """cuda""" or not is_xformers_available()			,							reason="""XFormers attention is only available with CUDA and `xformers` installed"""			,							)
		def 				__A							(  self					)							->  Dict:

					'''simple docstring'''
					self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3					)





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

					'''simple docstring'''
					self._test_inference_batch_single_identical(expected_max_diff=2E-3					)





		def 				__A							(  self					)							->  int:

					'''simple docstring'''
					__UpperCAmelCase :   Tuple    					=  self.get_dummy_components()
					__UpperCAmelCase :   int    					=  self.pipeline_class(**lowercase_					)
					pipe.to(lowercase_					)
					pipe.set_progress_bar_config(disable=lowercase_					)
					with tempfile.TemporaryDirectory() as tmpdir:
								try:
											# save_pretrained is not implemented for Multi-ControlNet
											pipe.save_pretrained(lowercase_					)
								except NotImplementedError:
											pass



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





		def 				__A							(  self					)							->  Tuple:

					'''simple docstring'''
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()





		def 				__A							(  self					)							->  Dict:

					'''simple docstring'''
					__UpperCAmelCase :   List[str]    					=  ControlNetModel.from_pretrained("""lllyasviel/sd-controlnet-canny"""					)

					__UpperCAmelCase :   str    					=  StableDiffusionControlNetInpaintPipeline.from_pretrained(
					    """runwayml/stable-diffusion-inpainting"""			,							safety_checker=lowercase_			,							controlnet=lowercase_					)
					pipe.enable_model_cpu_offload()
					pipe.set_progress_bar_config(disable=lowercase_					)

					__UpperCAmelCase :   Optional[Any]    					=  torch.Generator(device="""cpu"""					).manual_seed(0					)
					__UpperCAmelCase :   Tuple    					=  load_image(
					    """https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"""					).resize((512, 512)					)

					__UpperCAmelCase :   Tuple    					=  load_image(
					    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
					    """/stable_diffusion_inpaint/input_bench_mask.png"""					).resize((512, 512)					)

					__UpperCAmelCase :   Optional[int]    					=  'pitch black hole'

					__UpperCAmelCase :   str    					=  load_image(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"""					).resize((512, 512)					)

					__UpperCAmelCase :   int    					=  pipe(
					    lowercase_			,							image=lowercase_			,							mask_image=lowercase_			,							control_image=lowercase_			,							generator=lowercase_			,							output_type="""np"""			,							num_inference_steps=3			,							)

					__UpperCAmelCase :   Optional[int]    					=  output.images[0]

					assert image.shape == (512, 512, 3)

					__UpperCAmelCase :   str    					=  load_numpy(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/inpaint.npy"""					)

					assert np.abs(expected_image - image					).max() < 9E-2





		def 				__A							(  self					)							->  int:

					'''simple docstring'''
					__UpperCAmelCase :   Union[str, Any]    					=  ControlNetModel.from_pretrained("""lllyasviel/control_v11p_sd15_inpaint"""					)

					__UpperCAmelCase :   Tuple    					=  StableDiffusionControlNetInpaintPipeline.from_pretrained(
					    """runwayml/stable-diffusion-v1-5"""			,							safety_checker=lowercase_			,							controlnet=lowercase_					)
					__UpperCAmelCase :   List[Any]    					=  DDIMScheduler.from_config(pipe.scheduler.config					)
					pipe.enable_model_cpu_offload()
					pipe.set_progress_bar_config(disable=lowercase_					)

					__UpperCAmelCase :   List[Any]    					=  torch.Generator(device="""cpu"""					).manual_seed(33					)

					__UpperCAmelCase :   Dict    					=  load_image(
					    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png"""					)
					__UpperCAmelCase :   Any    					=  init_image.resize((512, 512)					)

					__UpperCAmelCase :   List[Any]    					=  load_image(
					    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png"""					)
					__UpperCAmelCase :   Tuple    					=  mask_image.resize((512, 512)					)

					__UpperCAmelCase :   Optional[int]    					=  'a handsome man with ray-ban sunglasses'

					def make_inpaint_condition(__UpperCAmelCase			,							__UpperCAmelCase					):
								__UpperCAmelCase :   List[Any]    					=  np.array(image.convert("""RGB"""					)					).astype(np.floataa					) / 255.0
								__UpperCAmelCase :   Tuple    					=  np.array(image_mask.convert("""L"""					)					).astype(np.floataa					) / 255.0

								assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
								__UpperCAmelCase :   Union[str, Any]    					=  -1.0  # set as masked pixel
								__UpperCAmelCase :   str    					=  np.expand_dims(lowercase_			,							0					).transpose(0			,							3			,							1			,							2					)
								__UpperCAmelCase :   Optional[Any]    					=  torch.from_numpy(lowercase_					)
								return image

					__UpperCAmelCase :   int    					=  make_inpaint_condition(lowercase_			,							lowercase_					)

					__UpperCAmelCase :   Dict    					=  pipe(
					    lowercase_			,							image=lowercase_			,							mask_image=lowercase_			,							control_image=lowercase_			,							guidance_scale=9.0			,							eta=1.0			,							generator=lowercase_			,							num_inference_steps=20			,							output_type="""np"""			,							)
					__UpperCAmelCase :   List[str]    					=  output.images[0]

					assert image.shape == (512, 512, 3)

					__UpperCAmelCase :   Optional[Any]    					=  load_numpy(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/boy_ray_ban.npy"""					)

					assert np.abs(expected_image - image					).max() < 9E-2












"""simple docstring"""


import importlib.metadata
import subprocess
import sys




def 					a			(  UpperCamelCase_				:					bool = False		)	->		Optional[Any]:
 __A						:							List[Any]      =						False
 if "IPython" in sys.modules:
  __A						:							List[Any]      =						'google.colab' in str(sys.modules['IPython'].get_ipython()		)

 if in_colab:
  if upgrade:
   __A						:							Dict      =						['pip', 'install', '-U', 'torch']
   subprocess.run(UpperCamelCase_  ,						check=UpperCamelCase_		)
  # get the current version of torch
  __A						:							int      =						importlib.metadata.version('torch'		)
  __A						:							int      =						torch_version[: torch_version.rindex('.'		)]
  __A						:							Any      =						F"https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-{torch_version_trunc}-cp37-cp37m-linux_x86_64.whl"
  __A						:							int      =						['pip', 'install', xla_wheel]
  subprocess.run(UpperCamelCase_  ,						check=UpperCamelCase_		)
 else:
  raise RuntimeError('`install_xla` utility works only on google colab.'		)





