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 SCREAMING_SNAKE_CASE				(     lowercase_:	nn.Module				):
       _UpperCamelCase     :       Any  						=  {}
       for name in unet.attn_processors.keys():
              _UpperCamelCase     :       Tuple  						=  None if name.endswith('attn1.processor'				) else unet.config.cross_attention_dim
              if name.startswith('mid_block'				):
                     _UpperCamelCase     :       Tuple  						=  unet.config.block_out_channels[-1]
              elif name.startswith('up_blocks'				):
                     _UpperCamelCase     :       Dict  						=  int(name[len('up_blocks.'				)]				)
                     _UpperCamelCase     :       Optional[int]  						=  list(reversed(unet.config.block_out_channels				)				)[block_id]
              elif name.startswith('down_blocks'				):
                     _UpperCamelCase     :       int  						=  int(name[len('down_blocks.'				)]				)
                     _UpperCamelCase     :       List[str]  						=  unet.config.block_out_channels[block_id]
              _UpperCamelCase     :       List[Any]  						=  (
                  LoRAAttnProcessor2_0 if hasattr(lowercase_		,						'scaled_dot_product_attention'				) else LoRAAttnProcessor
              )
              _UpperCamelCase     :       Dict  						=  lora_attn_processor_class(
                  hidden_size=lowercase_		,						cross_attention_dim=lowercase_				)
       _UpperCamelCase     :       int  						=  AttnProcsLayers(lowercase_				)
       return lora_attn_procs, unet_lora_layers


def SCREAMING_SNAKE_CASE				(     lowercase_:	nn.Module				):
       _UpperCamelCase     :       Union[str, Any]  						=  {}
       _UpperCamelCase     :       Union[str, Any]  						=  (
           LoRAAttnProcessor2_0 if hasattr(lowercase_		,						'scaled_dot_product_attention'				) else LoRAAttnProcessor
       )
       for name, module in text_encoder.named_modules():
              if name.endswith(lowercase_				):
                     _UpperCamelCase     :       str  						=  lora_attn_processor_class(
                         hidden_size=module.out_proj.out_features		,						cross_attention_dim=lowercase_				)
       return text_lora_attn_procs


def SCREAMING_SNAKE_CASE				(     lowercase_:	nn.Module				):
       _UpperCamelCase     :       Tuple  						=  create_text_encoder_lora_attn_procs(lowercase_				)
       _UpperCamelCase     :       Dict  						=  AttnProcsLayers(lowercase_				)
       return text_encoder_lora_layers


def SCREAMING_SNAKE_CASE				(     lowercase_:	Dict		,						lowercase_:	Any=False				):
       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'				):
                            _UpperCamelCase     :       List[str]  						=  (
                                torch.randn_like(layer_module.up.weight				)
                                if randn_weight
                                else torch.zeros_like(layer_module.up.weight				)
                            )
                            _UpperCamelCase     :       List[str]  						=  torch.nn.Parameter(lowercase_				)





class    __lowerCAmelCase  (       unittest.TestCase						):



 '''simple docstring'''

 def 					__snake_case					(			self       :		List[str]  ) ->							Any:
        torch.manual_seed(0  )
        _UpperCamelCase     :       Tuple  						=  UNetaDConditionModel(
            block_out_channels=(3_2, 6_4)       ,			layers_per_block=2       ,			sample_size=3_2       ,			in_channels=4       ,			out_channels=4       ,			down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')       ,			up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')       ,			cross_attention_dim=3_2       ,			)
        _UpperCamelCase     :       List[str]  						=  DDIMScheduler(
            beta_start=0.0_0_0_8_5       ,			beta_end=0.0_1_2       ,			beta_schedule='scaled_linear'       ,			clip_sample=lowerCamelCase       ,			set_alpha_to_one=lowerCamelCase       ,			steps_offset=1       ,			)
        torch.manual_seed(0  )
        _UpperCamelCase     :       Tuple  						=  AutoencoderKL(
            block_out_channels=[3_2, 6_4]       ,			in_channels=3       ,			out_channels=3       ,			down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']       ,			up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']       ,			latent_channels=4       ,			)
        _UpperCamelCase     :       Any  						=  CLIPTextConfig(
            bos_token_id=0       ,			eos_token_id=2       ,			hidden_size=3_2       ,			intermediate_size=3_7       ,			layer_norm_eps=1e-05       ,			num_attention_heads=4       ,			num_hidden_layers=5       ,			pad_token_id=1       ,			vocab_size=1_0_0_0       ,			)
        _UpperCamelCase     :       Optional[int]  						=  CLIPTextModel(lowerCamelCase  )
        _UpperCamelCase     :       List[Any]  						=  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )

        _UpperCamelCase						,      _UpperCamelCase     :       Dict  						=  create_unet_lora_layers(lowerCamelCase  )
        _UpperCamelCase     :       int  						=  create_text_encoder_lora_layers(lowerCamelCase  )

        _UpperCamelCase     :       Any  						=  {
            'unet': unet,
            'scheduler': scheduler,
            'vae': vae,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'safety_checker': None,
            'feature_extractor': None,
        }
        _UpperCamelCase     :       str  						=  {
            '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 					__snake_case					(			self       :		Union[str, Any]  ) ->							List[str]:
        _UpperCamelCase     :       Tuple  						=  1
        _UpperCamelCase     :       Tuple  						=  1_0
        _UpperCamelCase     :       Optional[Any]  						=  4
        _UpperCamelCase     :       List[str]  						=  (3_2, 3_2)

        _UpperCamelCase     :       List[Any]  						=  torch.manual_seed(0  )
        _UpperCamelCase     :       str  						=  floats_tensor((batch_size, num_channels) + sizes  )
        _UpperCamelCase     :       Any  						=  torch.randint(1       ,			lowerCamelCase       ,			size=(batch_size, sequence_length)       ,			generator=lowerCamelCase  )

        _UpperCamelCase     :       List[str]  						=  {
            '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 					__snake_case					(			self       :		List[str]  ) ->							Union[str, Any]:
        _UpperCamelCase     :       Union[str, Any]  						=  7_7

        _UpperCamelCase     :       Optional[Any]  						=  torch.randint(2       ,			5_6       ,			size=(1, max_seq_length)       ,			generator=torch.manual_seed(0  )  )

        _UpperCamelCase     :       List[str]  						=  {}
        _UpperCamelCase     :       Optional[Any]  						=  inputs
        return prepared_inputs

 def 					__snake_case					(			self       :		Tuple       ,			lowerCamelCase       :		Dict  ) ->							Optional[Any]:
        _UpperCamelCase						,      _UpperCamelCase     :       Tuple  						=  self.get_dummy_components()
        LoraLoaderMixin.save_lora_weights(
            save_directory=lowerCamelCase       ,			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(lowerCamelCase       ,			'pytorch_lora_weights.bin'  )  )  )

 def 					__snake_case					(			self       :		Tuple  ) ->							Union[str, Any]:
        _UpperCamelCase						,      _UpperCamelCase     :       Optional[int]  						=  self.get_dummy_components()
        _UpperCamelCase     :       Union[str, Any]  						=  StableDiffusionPipeline(**lowerCamelCase  )
        _UpperCamelCase     :       Optional[Any]  						=  sd_pipe.to(lowerCamelCase  )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase  )

        _UpperCamelCase						,      _UpperCamelCase						,      _UpperCamelCase     :       Dict  						=  self.get_dummy_inputs()

        _UpperCamelCase     :       Union[str, Any]  						=  sd_pipe(**lowerCamelCase  ).images
        _UpperCamelCase     :       Tuple  						=  original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=lowerCamelCase       ,			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(lowerCamelCase       ,			'pytorch_lora_weights.bin'  )  )  )
               sd_pipe.load_lora_weights(lowerCamelCase  )

        _UpperCamelCase     :       Any  						=  sd_pipe(**lowerCamelCase  ).images
        _UpperCamelCase     :       Tuple  						=  lora_images[0, -3:, -3:, -1]

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

 def 					__snake_case					(			self       :		List[Any]  ) ->							int:
        _UpperCamelCase						,      _UpperCamelCase     :       Tuple  						=  self.get_dummy_components()
        _UpperCamelCase     :       Optional[Any]  						=  StableDiffusionPipeline(**lowerCamelCase  )
        _UpperCamelCase     :       Dict  						=  sd_pipe.to(lowerCamelCase  )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase  )

        _UpperCamelCase						,      _UpperCamelCase						,      _UpperCamelCase     :       List[Any]  						=  self.get_dummy_inputs()

        _UpperCamelCase     :       Union[str, Any]  						=  sd_pipe(**lowerCamelCase  ).images
        _UpperCamelCase     :       int  						=  original_images[0, -3:, -3:, -1]

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

        _UpperCamelCase     :       Dict  						=  sd_pipe(**lowerCamelCase  ).images
        _UpperCamelCase     :       List[Any]  						=  lora_images[0, -3:, -3:, -1]

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

 def 					__snake_case					(			self       :		Union[str, Any]  ) ->							Dict:
        _UpperCamelCase						,      _UpperCamelCase     :       int  						=  self.get_dummy_components()
        _UpperCamelCase     :       Any  						=  lora_components['unet_lora_attn_procs']
        _UpperCamelCase     :       Any  						=  StableDiffusionPipeline(**lowerCamelCase  )
        _UpperCamelCase     :       str  						=  sd_pipe.to(lowerCamelCase  )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase  )

        _UpperCamelCase						,      _UpperCamelCase						,      _UpperCamelCase     :       List[Any]  						=  self.get_dummy_inputs()

        _UpperCamelCase     :       Optional[int]  						=  sd_pipe(**lowerCamelCase  ).images
        _UpperCamelCase     :       List[str]  						=  original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
               _UpperCamelCase     :       str  						=  sd_pipe.unet
               unet.set_attn_processor(lowerCamelCase  )
               unet.save_attn_procs(lowerCamelCase  )
               self.assertTrue(os.path.isfile(os.path.join(lowerCamelCase       ,			'pytorch_lora_weights.bin'  )  )  )
               sd_pipe.load_lora_weights(lowerCamelCase  )

        _UpperCamelCase     :       int  						=  sd_pipe(**lowerCamelCase  ).images
        _UpperCamelCase     :       Tuple  						=  lora_images[0, -3:, -3:, -1]

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

 def 					__snake_case					(			self       :		Union[str, Any]  ) ->							Any:
        _UpperCamelCase						,      _UpperCamelCase     :       List[Any]  						=  self.get_dummy_components()
        _UpperCamelCase     :       Optional[Any]  						=  StableDiffusionPipeline(**lowerCamelCase  )

        _UpperCamelCase     :       str  						=  self.get_dummy_tokens()

        # inference without lora
        _UpperCamelCase     :       Tuple  						=  pipe.text_encoder(**lowerCamelCase  )[0]
        assert outputs_without_lora.shape == (1, 7_7, 3_2)

        # create lora_attn_procs with zeroed out up.weights
        _UpperCamelCase     :       Dict  						=  create_text_encoder_lora_attn_procs(pipe.text_encoder  )
        set_lora_up_weights(lowerCamelCase       ,			randn_weight=lowerCamelCase  )

        # monkey patch
        pipe._modify_text_encoder(lowerCamelCase  )

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

        # inference with lora
        _UpperCamelCase     :       Any  						=  pipe.text_encoder(**lowerCamelCase  )[0]
        assert outputs_with_lora.shape == (1, 7_7, 3_2)

        assert torch.allclose(
            lowerCamelCase       ,			lowerCamelCase  ), "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
        _UpperCamelCase     :       Dict  						=  create_text_encoder_lora_attn_procs(pipe.text_encoder  )
        set_lora_up_weights(lowerCamelCase       ,			randn_weight=lowerCamelCase  )

        # monkey patch
        pipe._modify_text_encoder(lowerCamelCase  )

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

        # inference with lora
        _UpperCamelCase     :       Optional[int]  						=  pipe.text_encoder(**lowerCamelCase  )[0]
        assert outputs_with_lora.shape == (1, 7_7, 3_2)

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

 def 					__snake_case					(			self       :		Any  ) ->							Optional[int]:
        _UpperCamelCase						,      _UpperCamelCase     :       Optional[int]  						=  self.get_dummy_components()
        _UpperCamelCase     :       Union[str, Any]  						=  StableDiffusionPipeline(**lowerCamelCase  )

        _UpperCamelCase     :       Optional[Any]  						=  self.get_dummy_tokens()

        # inference without lora
        _UpperCamelCase     :       int  						=  pipe.text_encoder(**lowerCamelCase  )[0]
        assert outputs_without_lora.shape == (1, 7_7, 3_2)

        # create lora_attn_procs with randn up.weights
        _UpperCamelCase     :       Dict  						=  create_text_encoder_lora_attn_procs(pipe.text_encoder  )
        set_lora_up_weights(lowerCamelCase       ,			randn_weight=lowerCamelCase  )

        # monkey patch
        pipe._modify_text_encoder(lowerCamelCase  )

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

        # inference with lora
        _UpperCamelCase     :       str  						=  pipe.text_encoder(**lowerCamelCase  )[0]
        assert outputs_with_lora.shape == (1, 7_7, 3_2)

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

        # remove monkey patch
        pipe._remove_text_encoder_monkey_patch()

        # inference with removed lora
        _UpperCamelCase     :       Optional[Any]  						=  pipe.text_encoder(**lowerCamelCase  )[0]
        assert outputs_without_lora_removed.shape == (1, 7_7, 3_2)

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

 def 					__snake_case					(			self       :		Union[str, Any]  ) ->							Dict:
        _UpperCamelCase						,      _UpperCamelCase     :       Tuple  						=  self.get_dummy_components()
        _UpperCamelCase     :       Tuple  						=  StableDiffusionPipeline(**lowerCamelCase  )
        _UpperCamelCase     :       Dict  						=  sd_pipe.to(lowerCamelCase  )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase  )

        _UpperCamelCase						,      _UpperCamelCase						,      _UpperCamelCase     :       List[Any]  						=  self.get_dummy_inputs()

        with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=lowerCamelCase       ,			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(lowerCamelCase       ,			'pytorch_lora_weights.bin'  )  )  )
               sd_pipe.load_lora_weights(lowerCamelCase  )

        _UpperCamelCase     :       List[str]  						=  sd_pipe(**lowerCamelCase  ).images
        _UpperCamelCase     :       Optional[Any]  						=  lora_images[0, -3:, -3:, -1]

        _UpperCamelCase     :       Any  						=  sd_pipe(**lowerCamelCase       ,			cross_attention_kwargs={'scale': 0.5}  ).images
        _UpperCamelCase     :       List[str]  						=  lora_images_with_scale[0, -3:, -3:, -1]

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

 def 					__snake_case					(			self       :		Optional[Any]  ) ->							List[Any]:
        with tempfile.TemporaryDirectory() as tmpdirname:
               self.create_lora_weight_file(lowerCamelCase  )

               _UpperCamelCase						,      _UpperCamelCase     :       Dict  						=  self.get_dummy_components()
               _UpperCamelCase     :       Optional[Any]  						=  StableDiffusionPipeline(**lowerCamelCase  )
               _UpperCamelCase     :       Dict  						=  sd_pipe.to(lowerCamelCase  )
               sd_pipe.set_progress_bar_config(disable=lowerCamelCase  )

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

            # load LoRA weight file
               sd_pipe.load_lora_weights(lowerCamelCase  )

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

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

               _UpperCamelCase						,      _UpperCamelCase     :       Tuple  						=  self.get_dummy_components()
               _UpperCamelCase     :       Union[str, Any]  						=  StableDiffusionPipeline(**lowerCamelCase  )
               _UpperCamelCase     :       List[Any]  						=  sd_pipe.to(lowerCamelCase  )
               sd_pipe.set_progress_bar_config(disable=lowerCamelCase  )

               # 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(lowerCamelCase       ,			lowerCamelCase  ):
                             self.assertIsInstance(module.processor       ,			lowerCamelCase  )

            # load LoRA weight file
               sd_pipe.load_lora_weights(lowerCamelCase  )

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

 @unittest.skipIf(torch_device != 'cuda'       ,			'This test is supposed to run on GPU'  )
 def 					__snake_case					(			self       :		List[str]  ) ->							List[str]:
        _UpperCamelCase						,      _UpperCamelCase     :       str  						=  self.get_dummy_components()
        _UpperCamelCase     :       Any  						=  StableDiffusionPipeline(**lowerCamelCase  )
        _UpperCamelCase     :       Dict  						=  sd_pipe.to(lowerCamelCase  )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase  )

        _UpperCamelCase						,      _UpperCamelCase						,      _UpperCamelCase     :       Tuple  						=  self.get_dummy_inputs()

        # enable XFormers
        sd_pipe.enable_xformers_memory_efficient_attention()

        _UpperCamelCase     :       Optional[int]  						=  sd_pipe(**lowerCamelCase  ).images
        _UpperCamelCase     :       str  						=  original_images[0, -3:, -3:, -1]

        with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=lowerCamelCase       ,			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(lowerCamelCase       ,			'pytorch_lora_weights.bin'  )  )  )
               sd_pipe.load_lora_weights(lowerCamelCase  )

        _UpperCamelCase     :       List[str]  						=  sd_pipe(**lowerCamelCase  ).images
        _UpperCamelCase     :       int  						=  lora_images[0, -3:, -3:, -1]

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




from .chunk_utils import chunk_layer
from .data_transforms import make_atomaa_masks
from .feats import atomaa_to_atomaa, frames_and_literature_positions_to_atomaa_pos, torsion_angles_to_frames
from .loss import compute_predicted_aligned_error, compute_tm
from .protein import Protein as OFProtein
from .protein import to_pdb
from .rigid_utils import Rigid, Rotation
from .tensor_utils import dict_multimap, flatten_final_dims, permute_final_dims



