
"""simple docstring"""


import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device



def 					_lowerCamelCase(__UpperCamelCase )       ->				Any:
	_lowerCAmelCase			={}
	for name in unet.attn_processors.keys():
		_lowerCAmelCase			=None if name.endswith("""attn1.processor""" ) else unet.config.cross_attention_dim
		if name.startswith("""mid_block""" ):
			_lowerCAmelCase			=unet.config.block_out_channels[-1]
		elif name.startswith("""up_blocks""" ):
			_lowerCAmelCase			=int(name[len("""up_blocks.""" )] )
			_lowerCAmelCase			=list(reversed(unet.config.block_out_channels ) )[block_id]
		elif name.startswith("""down_blocks""" ):
			_lowerCAmelCase			=int(name[len("""down_blocks.""" )] )
			_lowerCAmelCase			=unet.config.block_out_channels[block_id]
		_lowerCAmelCase			=(
		    LoRAAttnProcessor2_0 if hasattr(lowercase_					,				"""scaled_dot_product_attention""" ) else LoRAAttnProcessor
		)
		_lowerCAmelCase			=lora_attn_processor_class(
		    hidden_size=lowercase_					,				cross_attention_dim=lowercase_ )
	_lowerCAmelCase			=AttnProcsLayers(lowercase_ )
	return lora_attn_procs, unet_lora_layers



def 					_lowerCamelCase(__UpperCamelCase )       ->				str:
	_lowerCAmelCase			={}
	_lowerCAmelCase			=(
	    LoRAAttnProcessor2_0 if hasattr(lowercase_					,				"""scaled_dot_product_attention""" ) else LoRAAttnProcessor
	)
	for name, module in text_encoder.named_modules():
		if name.endswith(lowercase_ ):
			_lowerCAmelCase			=lora_attn_processor_class(
			    hidden_size=module.out_proj.out_features					,				cross_attention_dim=lowercase_ )
	return text_lora_attn_procs



def 					_lowerCamelCase(__UpperCamelCase )       ->				Union[str, Any]:
	_lowerCAmelCase			=create_text_encoder_lora_attn_procs(lowercase_ )
	_lowerCAmelCase			=AttnProcsLayers(lowercase_ )
	return text_encoder_lora_layers



def 					_lowerCamelCase(__UpperCamelCase					,				__UpperCamelCase=False )       ->				int:
	for _, attn_proc in text_lora_attn_procs.items():
		# set up.weights
		for layer_name, layer_module in attn_proc.named_modules():
			if layer_name.endswith("""_lora""" ):
				_lowerCAmelCase			=(
				    torch.randn_like(layer_module.up.weight )
				    if randn_weight
				    else torch.zeros_like(layer_module.up.weight )
				)
				_lowerCAmelCase			=torch.nn.Parameter(lowercase_ )







class 	lowerCamelCase__      (      unittest.TestCase       ):





			'''simple docstring'''



			def       _lowerCAmelCase	(							self						)				->  Any:
				torch.manual_seed(0						)
				_lowerCAmelCase			=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  ,  )
				_lowerCAmelCase			=DDIMScheduler(
				    beta_start=0.0_0_0_8_5  ,  beta_end=0.0_1_2  ,  beta_schedule="""scaled_linear"""  ,  clip_sample=__UpperCAmelCase  ,  set_alpha_to_one=__UpperCAmelCase  ,  steps_offset=1  ,  )
				torch.manual_seed(0						)
				_lowerCAmelCase			=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  ,  )
				_lowerCAmelCase			=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=10_00  ,  )
				_lowerCAmelCase			=CLIPTextModel(__UpperCAmelCase						)
				_lowerCAmelCase			=CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

				_lowerCAmelCase			=create_unet_lora_layers(__UpperCAmelCase						)
				_lowerCAmelCase			=create_text_encoder_lora_layers(__UpperCAmelCase						)

				_lowerCAmelCase			={
				    'unet': unet,
				    'scheduler': scheduler,
				    'vae': vae,
				    'text_encoder': text_encoder,
				    'tokenizer': tokenizer,
				    'safety_checker': None,
				    'feature_extractor': None,
				}
				_lowerCAmelCase			={
				    'unet_lora_layers': unet_lora_layers,
				    'text_encoder_lora_layers': text_encoder_lora_layers,
				    'unet_lora_attn_procs': unet_lora_attn_procs,
				}
				return pipeline_components, lora_components



			def       _lowerCAmelCase	(							self						)				->  List[str]:
				_lowerCAmelCase			=1
				_lowerCAmelCase			=10
				_lowerCAmelCase			=4
				_lowerCAmelCase			=(32, 32)

				_lowerCAmelCase			=torch.manual_seed(0						)
				_lowerCAmelCase			=floats_tensor((batch_size, num_channels) + sizes						)
				_lowerCAmelCase			=torch.randint(1  ,  __UpperCAmelCase  ,  size=(batch_size, sequence_length)  ,  generator=__UpperCAmelCase						)

				_lowerCAmelCase			={
				    'prompt': 'A painting of a squirrel eating a burger',
				    'generator': generator,
				    'num_inference_steps': 2,
				    'guidance_scale': 6.0,
				    'output_type': 'numpy',
				}

				return noise, input_ids, pipeline_inputs

				# copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb



			def       _lowerCAmelCase	(							self						)				->  Union[str, Any]:
				_lowerCAmelCase			=77

				_lowerCAmelCase			=torch.randint(2  ,  56  ,  size=(1, max_seq_length)  ,  generator=torch.manual_seed(0						)						)

				_lowerCAmelCase			={}
				_lowerCAmelCase			=inputs
				return prepared_inputs



			def       _lowerCAmelCase	(							self  ,  __UpperCAmelCase						)				->  Optional[Any]:
				_lowerCAmelCase			=self.get_dummy_components()
				LoraLoaderMixin.save_lora_weights(
				    save_directory=__UpperCAmelCase  ,  unet_lora_layers=lora_components["""unet_lora_layers"""]  ,  text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]  ,  )
				self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase  ,  """pytorch_lora_weights.bin"""						)						)						)



			def       _lowerCAmelCase	(							self						)				->  Union[str, Any]:
				_lowerCAmelCase			=self.get_dummy_components()
				_lowerCAmelCase			=StableDiffusionPipeline(**__UpperCAmelCase						)
				_lowerCAmelCase			=sd_pipe.to(__UpperCAmelCase						)
				sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

				_lowerCAmelCase			=self.get_dummy_inputs()

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase						).images
				_lowerCAmelCase			=original_images[0, -3:, -3:, -1]

				with tempfile.TemporaryDirectory() as tmpdirname:
					LoraLoaderMixin.save_lora_weights(
					    save_directory=__UpperCAmelCase  ,  unet_lora_layers=lora_components["""unet_lora_layers"""]  ,  text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]  ,  )
					self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase  ,  """pytorch_lora_weights.bin"""						)						)						)
					sd_pipe.load_lora_weights(__UpperCAmelCase						)

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase						).images
				_lowerCAmelCase			=lora_images[0, -3:, -3:, -1]

				# Outputs shouldn't match.
				self.assertFalse(torch.allclose(torch.from_numpy(__UpperCAmelCase						)  ,  torch.from_numpy(__UpperCAmelCase						)						)						)



			def       _lowerCAmelCase	(							self						)				->  int:
				_lowerCAmelCase			=self.get_dummy_components()
				_lowerCAmelCase			=StableDiffusionPipeline(**__UpperCAmelCase						)
				_lowerCAmelCase			=sd_pipe.to(__UpperCAmelCase						)
				sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

				_lowerCAmelCase			=self.get_dummy_inputs()

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase						).images
				_lowerCAmelCase			=original_images[0, -3:, -3:, -1]

				with tempfile.TemporaryDirectory() as tmpdirname:
					LoraLoaderMixin.save_lora_weights(
					    save_directory=__UpperCAmelCase  ,  unet_lora_layers=lora_components["""unet_lora_layers"""]  ,  text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]  ,  safe_serialization=__UpperCAmelCase  ,  )
					self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase  ,  """pytorch_lora_weights.safetensors"""						)						)						)
					sd_pipe.load_lora_weights(__UpperCAmelCase						)

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase						).images
				_lowerCAmelCase			=lora_images[0, -3:, -3:, -1]

				# Outputs shouldn't match.
				self.assertFalse(torch.allclose(torch.from_numpy(__UpperCAmelCase						)  ,  torch.from_numpy(__UpperCAmelCase						)						)						)



			def       _lowerCAmelCase	(							self						)				->  Dict:
				_lowerCAmelCase			=self.get_dummy_components()
				_lowerCAmelCase			=lora_components['unet_lora_attn_procs']
				_lowerCAmelCase			=StableDiffusionPipeline(**__UpperCAmelCase						)
				_lowerCAmelCase			=sd_pipe.to(__UpperCAmelCase						)
				sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

				_lowerCAmelCase			=self.get_dummy_inputs()

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase						).images
				_lowerCAmelCase			=original_images[0, -3:, -3:, -1]

				with tempfile.TemporaryDirectory() as tmpdirname:
					_lowerCAmelCase			=sd_pipe.unet
					unet.set_attn_processor(__UpperCAmelCase						)
					unet.save_attn_procs(__UpperCAmelCase						)
					self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase  ,  """pytorch_lora_weights.bin"""						)						)						)
					sd_pipe.load_lora_weights(__UpperCAmelCase						)

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase						).images
				_lowerCAmelCase			=lora_images[0, -3:, -3:, -1]

				# Outputs shouldn't match.
				self.assertFalse(torch.allclose(torch.from_numpy(__UpperCAmelCase						)  ,  torch.from_numpy(__UpperCAmelCase						)						)						)



			def       _lowerCAmelCase	(							self						)				->  Any:
				_lowerCAmelCase			=self.get_dummy_components()
				_lowerCAmelCase			=StableDiffusionPipeline(**__UpperCAmelCase						)

				_lowerCAmelCase			=self.get_dummy_tokens()

				# inference without lora
				_lowerCAmelCase			=pipe.text_encoder(**__UpperCAmelCase						)[0]
				assert outputs_without_lora.shape == (1, 77, 32)

				# create lora_attn_procs with zeroed out up.weights
				_lowerCAmelCase			=create_text_encoder_lora_attn_procs(pipe.text_encoder						)
				set_lora_up_weights(__UpperCAmelCase  ,  randn_weight=__UpperCAmelCase						)

				# monkey patch
				pipe._modify_text_encoder(__UpperCAmelCase						)

				# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
				del text_attn_procs
				gc.collect()

				# inference with lora
				_lowerCAmelCase			=pipe.text_encoder(**__UpperCAmelCase						)[0]
				assert outputs_with_lora.shape == (1, 77, 32)

				assert torch.allclose(
				    __UpperCAmelCase  ,  __UpperCAmelCase						), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

				# create lora_attn_procs with randn up.weights
				_lowerCAmelCase			=create_text_encoder_lora_attn_procs(pipe.text_encoder						)
				set_lora_up_weights(__UpperCAmelCase  ,  randn_weight=__UpperCAmelCase						)

				# monkey patch
				pipe._modify_text_encoder(__UpperCAmelCase						)

				# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
				del text_attn_procs
				gc.collect()

				# inference with lora
				_lowerCAmelCase			=pipe.text_encoder(**__UpperCAmelCase						)[0]
				assert outputs_with_lora.shape == (1, 77, 32)

				assert not torch.allclose(
				    __UpperCAmelCase  ,  __UpperCAmelCase						), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"



			def       _lowerCAmelCase	(							self						)				->  Optional[int]:
				_lowerCAmelCase			=self.get_dummy_components()
				_lowerCAmelCase			=StableDiffusionPipeline(**__UpperCAmelCase						)

				_lowerCAmelCase			=self.get_dummy_tokens()

				# inference without lora
				_lowerCAmelCase			=pipe.text_encoder(**__UpperCAmelCase						)[0]
				assert outputs_without_lora.shape == (1, 77, 32)

				# create lora_attn_procs with randn up.weights
				_lowerCAmelCase			=create_text_encoder_lora_attn_procs(pipe.text_encoder						)
				set_lora_up_weights(__UpperCAmelCase  ,  randn_weight=__UpperCAmelCase						)

				# monkey patch
				pipe._modify_text_encoder(__UpperCAmelCase						)

				# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
				del text_attn_procs
				gc.collect()

				# inference with lora
				_lowerCAmelCase			=pipe.text_encoder(**__UpperCAmelCase						)[0]
				assert outputs_with_lora.shape == (1, 77, 32)

				assert not torch.allclose(
				    __UpperCAmelCase  ,  __UpperCAmelCase						), "lora outputs should be different to without lora outputs"

				# remove monkey patch
				pipe._remove_text_encoder_monkey_patch()

				# inference with removed lora
				_lowerCAmelCase			=pipe.text_encoder(**__UpperCAmelCase						)[0]
				assert outputs_without_lora_removed.shape == (1, 77, 32)

				assert torch.allclose(
				    __UpperCAmelCase  ,  __UpperCAmelCase						), "remove lora monkey patch should restore the original outputs"



			def       _lowerCAmelCase	(							self						)				->  Dict:
				_lowerCAmelCase			=self.get_dummy_components()
				_lowerCAmelCase			=StableDiffusionPipeline(**__UpperCAmelCase						)
				_lowerCAmelCase			=sd_pipe.to(__UpperCAmelCase						)
				sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

				_lowerCAmelCase			=self.get_dummy_inputs()

				with tempfile.TemporaryDirectory() as tmpdirname:
					LoraLoaderMixin.save_lora_weights(
					    save_directory=__UpperCAmelCase  ,  unet_lora_layers=lora_components["""unet_lora_layers"""]  ,  text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]  ,  )
					self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase  ,  """pytorch_lora_weights.bin"""						)						)						)
					sd_pipe.load_lora_weights(__UpperCAmelCase						)

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase						).images
				_lowerCAmelCase			=lora_images[0, -3:, -3:, -1]

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase  ,  cross_attention_kwargs={"""scale""": 0.5}						).images
				_lowerCAmelCase			=lora_images_with_scale[0, -3:, -3:, -1]

				# Outputs shouldn't match.
				self.assertFalse(
				    torch.allclose(torch.from_numpy(__UpperCAmelCase						)  ,  torch.from_numpy(__UpperCAmelCase						)						)						)



			def       _lowerCAmelCase	(							self						)				->  List[Any]:
				with tempfile.TemporaryDirectory() as tmpdirname:
					self.create_lora_weight_file(__UpperCAmelCase						)

					_lowerCAmelCase			=self.get_dummy_components()
					_lowerCAmelCase			=StableDiffusionPipeline(**__UpperCAmelCase						)
					_lowerCAmelCase			=sd_pipe.to(__UpperCAmelCase						)
					sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

					# check if vanilla attention processors are used
					for _, module in sd_pipe.unet.named_modules():
						if isinstance(__UpperCAmelCase  ,  __UpperCAmelCase						):
							self.assertIsInstance(module.processor  ,  (AttnProcessor, AttnProcessor2_0)						)

            # load LoRA weight file
					sd_pipe.load_lora_weights(__UpperCAmelCase						)

					# check if lora attention processors are used
					for _, module in sd_pipe.unet.named_modules():
						if isinstance(__UpperCAmelCase  ,  __UpperCAmelCase						):
							_lowerCAmelCase			=(
							    LoRAAttnProcessor2_0 if hasattr(__UpperCAmelCase  ,  """scaled_dot_product_attention"""						) else LoRAAttnProcessor
							)
							self.assertIsInstance(module.processor  ,  __UpperCAmelCase						)



			@unittest.skipIf(torch_device != """cuda"""  ,  """This test is supposed to run on GPU"""						)
			def       _lowerCAmelCase	(							self						)				->  Optional[Any]:
				with tempfile.TemporaryDirectory() as tmpdirname:
					self.create_lora_weight_file(__UpperCAmelCase						)

					_lowerCAmelCase			=self.get_dummy_components()
					_lowerCAmelCase			=StableDiffusionPipeline(**__UpperCAmelCase						)
					_lowerCAmelCase			=sd_pipe.to(__UpperCAmelCase						)
					sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

					# enable XFormers
					sd_pipe.enable_xformers_memory_efficient_attention()

					# check if xFormers attention processors are used
					for _, module in sd_pipe.unet.named_modules():
						if isinstance(__UpperCAmelCase  ,  __UpperCAmelCase						):
							self.assertIsInstance(module.processor  ,  __UpperCAmelCase						)

            # load LoRA weight file
					sd_pipe.load_lora_weights(__UpperCAmelCase						)

					# check if lora attention processors are used
					for _, module in sd_pipe.unet.named_modules():
						if isinstance(__UpperCAmelCase  ,  __UpperCAmelCase						):
							self.assertIsInstance(module.processor  ,  __UpperCAmelCase						)



			@unittest.skipIf(torch_device != """cuda"""  ,  """This test is supposed to run on GPU"""						)
			def       _lowerCAmelCase	(							self						)				->  List[str]:
				_lowerCAmelCase			=self.get_dummy_components()
				_lowerCAmelCase			=StableDiffusionPipeline(**__UpperCAmelCase						)
				_lowerCAmelCase			=sd_pipe.to(__UpperCAmelCase						)
				sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

				_lowerCAmelCase			=self.get_dummy_inputs()

				# enable XFormers
				sd_pipe.enable_xformers_memory_efficient_attention()

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase						).images
				_lowerCAmelCase			=original_images[0, -3:, -3:, -1]

				with tempfile.TemporaryDirectory() as tmpdirname:
					LoraLoaderMixin.save_lora_weights(
					    save_directory=__UpperCAmelCase  ,  unet_lora_layers=lora_components["""unet_lora_layers"""]  ,  text_encoder_lora_layers=lora_components["""text_encoder_lora_layers"""]  ,  )
					self.assertTrue(os.path.isfile(os.path.join(__UpperCAmelCase  ,  """pytorch_lora_weights.bin"""						)						)						)
					sd_pipe.load_lora_weights(__UpperCAmelCase						)

				_lowerCAmelCase			=sd_pipe(**__UpperCAmelCase						).images
				_lowerCAmelCase			=lora_images[0, -3:, -3:, -1]

				# Outputs shouldn't match.
				self.assertFalse(torch.allclose(torch.from_numpy(__UpperCAmelCase						)  ,  torch.from_numpy(__UpperCAmelCase						)						)						)

"""simple docstring"""


from typing import Dict, Optional, Tuple

import torch



def 					_lowerCamelCase(__UpperCamelCase )       ->				torch.Tensor:
	_lowerCAmelCase			=boundaries[1] - boundaries[0]
	_lowerCAmelCase			=boundaries + step / 2
	_lowerCAmelCase			=torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1 )]					,				dim=0 )
	return bin_centers



def 					_lowerCamelCase(__UpperCamelCase					,				__UpperCamelCase					,				)       ->				Tuple[torch.Tensor, torch.Tensor]:
	_lowerCAmelCase			=_calculate_bin_centers(__UpperCamelCase )
	return (
	    torch.sum(aligned_distance_error_probs * bin_centers					,				dim=-1 ),
	    bin_centers[-1],
	)



def 					_lowerCamelCase(__UpperCamelCase					,				__UpperCamelCase = 31					,				__UpperCamelCase = 64					,				**__UpperCamelCase					,				)       ->				Dict[str, torch.Tensor]:
	_lowerCAmelCase			=torch.linspace(0					,				__UpperCamelCase					,				steps=(no_bins - 1)					,				device=logits.device )

	_lowerCAmelCase			=torch.nn.functional.softmax(__UpperCamelCase					,				dim=-1 )
	_lowerCAmelCase		,							_lowerCAmelCase			=_calculate_expected_aligned_error(
	    alignment_confidence_breaks=__UpperCamelCase					,				aligned_distance_error_probs=__UpperCamelCase					,				)

	return {
	    "aligned_confidence_probs": aligned_confidence_probs,
	    "predicted_aligned_error": predicted_aligned_error,
	    "max_predicted_aligned_error": max_predicted_aligned_error,
	}



def 					_lowerCamelCase(__UpperCamelCase					,				__UpperCamelCase = None					,				__UpperCamelCase = 31					,				__UpperCamelCase = 64					,				__UpperCamelCase = 1E-8					,				**__UpperCamelCase					,				)       ->				torch.Tensor:
	if residue_weights is None:
		_lowerCAmelCase			=logits.new_ones(logits.shape[-2] )

	_lowerCAmelCase			=torch.linspace(0					,				__UpperCamelCase					,				steps=(no_bins - 1)					,				device=logits.device )

	_lowerCAmelCase			=_calculate_bin_centers(__UpperCamelCase )
	torch.sum(__UpperCamelCase )
	_lowerCAmelCase			=logits.shape[-2]
	_lowerCAmelCase			=max(__UpperCamelCase					,				19 )

	_lowerCAmelCase			=1.24 * (clipped_n - 15) ** (1.0 / 3) - 1.8

	_lowerCAmelCase			=torch.nn.functional.softmax(__UpperCamelCase					,				dim=-1 )

	_lowerCAmelCase			=1.0 / (1 + (bin_centers**2) / (da**2))
	_lowerCAmelCase			=torch.sum(probs * tm_per_bin					,				dim=-1 )

	_lowerCAmelCase			=residue_weights / (eps + residue_weights.sum())
	_lowerCAmelCase			=torch.sum(predicted_tm_term * normed_residue_mask					,				dim=-1 )

	_lowerCAmelCase			=per_alignment * residue_weights

	_lowerCAmelCase			=(weighted == torch.max(__UpperCamelCase )).nonzero()[0]
	return per_alignment[tuple(__UpperCamelCase )]



