



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

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()



class    SCREAMING_SNAKE_CASE__						(      UpperCamelCase_   ,    unittest.TestCase				):
			__lowerCAmelCase						:     Tuple		=		UnCLIPImageVariationPipeline
			__lowerCAmelCase						:     str		=		IMAGE_VARIATION_PARAMS - {'height', 'width', 'guidance_scale'}
			__lowerCAmelCase						:     Dict		=		IMAGE_VARIATION_BATCH_PARAMS

			__lowerCAmelCase						:     Tuple		=		[
			    'generator',
			    'return_dict',
			    'decoder_num_inference_steps',
			    'super_res_num_inference_steps',
			]
			__lowerCAmelCase						:     Optional[int]		=		False

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Union[str, Any]:





				'''simple docstring'''




				return 32

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[Any]:





				'''simple docstring'''




				return 32

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[int]:





				'''simple docstring'''




				return self.time_input_dim

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Tuple:





				'''simple docstring'''




				return self.time_input_dim * 4

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Tuple:





				'''simple docstring'''




				return 100

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   int:





				'''simple docstring'''




				UpperCAmelCase			:      int			=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""		)
				return tokenizer

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Dict:





				'''simple docstring'''




				torch.manual_seed(0		)
				UpperCAmelCase			:      List[Any]			=			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=1000   ,  )
				return CLIPTextModelWithProjection(_SCREAMING_SNAKE_CASE		)

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   int:





				'''simple docstring'''




				torch.manual_seed(0		)
				UpperCAmelCase			:      str			=			CLIPVisionConfig(
				    hidden_size=self.text_embedder_hidden_size   ,  projection_dim=self.text_embedder_hidden_size   ,  num_hidden_layers=5   ,  num_attention_heads=4   ,  image_size=32   ,  intermediate_size=37   ,  patch_size=1   ,  )
				return CLIPVisionModelWithProjection(_SCREAMING_SNAKE_CASE		)

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[int]:





				'''simple docstring'''




				torch.manual_seed(0		)

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

				UpperCAmelCase			:      Union[str, Any]			=			UnCLIPTextProjModel(**_SCREAMING_SNAKE_CASE		)
				return model

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Tuple:





				'''simple docstring'''




				torch.manual_seed(0		)

				UpperCAmelCase			:      str			=			{
				    '''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''',
				}

				UpperCAmelCase			:      Union[str, Any]			=			UNetaDConditionModel(**_SCREAMING_SNAKE_CASE		)
				return model

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[int]:





				'''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 			SCREAMING_SNAKE_CASE  (     self		)						->   List[Any]:





				'''simple docstring'''




				torch.manual_seed(0		)

				UpperCAmelCase			:      Any			=			UNetaDModel(**self.dummy_super_res_kwargs		)
				return model

			@property
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Dict:





				'''simple docstring'''




				torch.manual_seed(1		)

				UpperCAmelCase			:      Union[str, Any]			=			UNetaDModel(**self.dummy_super_res_kwargs		)
				return model

			def 			SCREAMING_SNAKE_CASE  (     self		)						->   str:





				'''simple docstring'''




				UpperCAmelCase			:      Optional[Any]			=			self.dummy_decoder
				UpperCAmelCase			:      Dict			=			self.dummy_text_proj
				UpperCAmelCase			:      Any			=			self.dummy_text_encoder
				UpperCAmelCase			:      Any			=			self.dummy_tokenizer
				UpperCAmelCase			:      Union[str, Any]			=			self.dummy_super_res_first
				UpperCAmelCase			:      int			=			self.dummy_super_res_last

				UpperCAmelCase			:      Any			=			UnCLIPScheduler(
				    variance_type="""learned_range"""   ,  prediction_type="""epsilon"""   ,  num_train_timesteps=1000   ,  )

				UpperCAmelCase			:      int			=			UnCLIPScheduler(
				    variance_type="""fixed_small_log"""   ,  prediction_type="""epsilon"""   ,  num_train_timesteps=1000   ,  )

				UpperCAmelCase			:      Dict			=			CLIPImageProcessor(crop_size=32   ,  size=32		)

				UpperCAmelCase			:      List[str]			=			self.dummy_image_encoder

				return {
				    "decoder": decoder,
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "text_proj": text_proj,
				    "feature_extractor": feature_extractor,
				    "image_encoder": image_encoder,
				    "super_res_first": super_res_first,
				    "super_res_last": super_res_last,
				    "decoder_scheduler": decoder_scheduler,
				    "super_res_scheduler": super_res_scheduler,
				}

			def 			SCREAMING_SNAKE_CASE  (     self   ,  _SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE=0   ,  _SCREAMING_SNAKE_CASE=True		)						->   Tuple:





				'''simple docstring'''




				UpperCAmelCase			:      Tuple			=			floats_tensor((1, 3, 32, 32)   ,  rng=random.Random(_SCREAMING_SNAKE_CASE		)		).to(_SCREAMING_SNAKE_CASE		)
				if str(_SCREAMING_SNAKE_CASE		).startswith("""mps"""		):
					UpperCAmelCase			:      Optional[Any]			=			torch.manual_seed(_SCREAMING_SNAKE_CASE		)
				else:
					UpperCAmelCase			:      Tuple			=			torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(_SCREAMING_SNAKE_CASE		)

				if pil_image:
					UpperCAmelCase			:      Optional[int]			=			input_image * 0.5 + 0.5
					UpperCAmelCase			:      Optional[int]			=			input_image.clamp(0   ,  1		)
					UpperCAmelCase			:      int			=			input_image.cpu().permute(0   ,  2   ,  3   ,  1		).float().numpy()
					UpperCAmelCase			:      Union[str, Any]			=			DiffusionPipeline.numpy_to_pil(_SCREAMING_SNAKE_CASE		)[0]

				return {
				    "image": input_image,
				    "generator": generator,
				    "decoder_num_inference_steps": 2,
				    "super_res_num_inference_steps": 2,
				    "output_type": "np",
				}

			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Dict:





				'''simple docstring'''




				UpperCAmelCase			:      Optional[Any]			=			'''cpu'''

				UpperCAmelCase			:      Any			=			self.get_dummy_components()

				UpperCAmelCase			:      Union[str, Any]			=			self.pipeline_class(**_SCREAMING_SNAKE_CASE		)
				UpperCAmelCase			:      str			=			pipe.to(_SCREAMING_SNAKE_CASE		)

				pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      int			=			self.get_dummy_inputs(_SCREAMING_SNAKE_CASE   ,  pil_image=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      Optional[int]			=			pipe(**_SCREAMING_SNAKE_CASE		)
				UpperCAmelCase			:      Tuple			=			output.images

				UpperCAmelCase			:      Optional[int]			=			self.get_dummy_inputs(_SCREAMING_SNAKE_CASE   ,  pil_image=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      Optional[Any]			=			pipe(
				    **_SCREAMING_SNAKE_CASE   ,  return_dict=_SCREAMING_SNAKE_CASE   ,  )[0]

				UpperCAmelCase			:      Any			=			image[0, -3:, -3:, -1]
				UpperCAmelCase			:      List[str]			=			image_from_tuple[0, -3:, -3:, -1]

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

				UpperCAmelCase			:      Union[str, Any]			=			np.array(
				    [
				        0.9997,
				        0.0002,
				        0.9997,
				        0.9997,
				        0.9969,
				        0.0023,
				        0.9997,
				        0.9969,
				        0.9970,
				    ]		)

				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 			SCREAMING_SNAKE_CASE  (     self		)						->   Union[str, Any]:





				'''simple docstring'''




				UpperCAmelCase			:      Any			=			'''cpu'''

				UpperCAmelCase			:      Union[str, Any]			=			self.get_dummy_components()

				UpperCAmelCase			:      Union[str, Any]			=			self.pipeline_class(**_SCREAMING_SNAKE_CASE		)
				UpperCAmelCase			:      Dict			=			pipe.to(_SCREAMING_SNAKE_CASE		)

				pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      List[Any]			=			self.get_dummy_inputs(_SCREAMING_SNAKE_CASE   ,  pil_image=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      int			=			pipe(**_SCREAMING_SNAKE_CASE		)
				UpperCAmelCase			:      Dict			=			output.images

				UpperCAmelCase			:      List[str]			=			self.get_dummy_inputs(_SCREAMING_SNAKE_CASE   ,  pil_image=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      Any			=			pipe(
				    **_SCREAMING_SNAKE_CASE   ,  return_dict=_SCREAMING_SNAKE_CASE   ,  )[0]

				UpperCAmelCase			:      Optional[int]			=			image[0, -3:, -3:, -1]
				UpperCAmelCase			:      Union[str, Any]			=			image_from_tuple[0, -3:, -3:, -1]

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

				UpperCAmelCase			:      str			=			np.array([0.9997, 0.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971]		)

				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 			SCREAMING_SNAKE_CASE  (     self		)						->   List[str]:





				'''simple docstring'''




				UpperCAmelCase			:      Any			=			'''cpu'''

				UpperCAmelCase			:      Optional[int]			=			self.get_dummy_components()

				UpperCAmelCase			:      Union[str, Any]			=			self.pipeline_class(**_SCREAMING_SNAKE_CASE		)
				UpperCAmelCase			:      int			=			pipe.to(_SCREAMING_SNAKE_CASE		)

				pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      Optional[Any]			=			self.get_dummy_inputs(_SCREAMING_SNAKE_CASE   ,  pil_image=_SCREAMING_SNAKE_CASE		)
				UpperCAmelCase			:      Optional[int]			=			[
				    pipeline_inputs['''image'''],
				    pipeline_inputs['''image'''],
				]

				UpperCAmelCase			:      Any			=			pipe(**_SCREAMING_SNAKE_CASE		)
				UpperCAmelCase			:      Tuple			=			output.images

				UpperCAmelCase			:      Tuple			=			self.get_dummy_inputs(_SCREAMING_SNAKE_CASE   ,  pil_image=_SCREAMING_SNAKE_CASE		)
				UpperCAmelCase			:      Union[str, Any]			=			[
				    tuple_pipeline_inputs['''image'''],
				    tuple_pipeline_inputs['''image'''],
				]

				UpperCAmelCase			:      Optional[Any]			=			pipe(
				    **_SCREAMING_SNAKE_CASE   ,  return_dict=_SCREAMING_SNAKE_CASE   ,  )[0]

				UpperCAmelCase			:      Any			=			image[0, -3:, -3:, -1]
				UpperCAmelCase			:      Dict			=			image_from_tuple[0, -3:, -3:, -1]

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

				UpperCAmelCase			:      str			=			np.array(
				    [
				        0.9997,
				        0.9989,
				        0.0008,
				        0.0021,
				        0.9960,
				        0.0018,
				        0.0014,
				        0.0002,
				        0.9933,
				    ]		)

				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 			SCREAMING_SNAKE_CASE  (     self		)						->   List[str]:





				'''simple docstring'''




				UpperCAmelCase			:      Optional[int]			=			torch.device("""cpu"""		)



				class    SCREAMING_SNAKE_CASE__						:
							__lowerCAmelCase						:     str		=		1
				UpperCAmelCase			:      Optional[int]			=			self.get_dummy_components()

				UpperCAmelCase			:      int			=			self.pipeline_class(**_SCREAMING_SNAKE_CASE		)
				UpperCAmelCase			:      List[Any]			=			pipe.to(_SCREAMING_SNAKE_CASE		)

				pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      Any			=			torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(0		)
				UpperCAmelCase			:      Any			=			pipe.decoder.dtype
				UpperCAmelCase			:      Dict			=			1

				UpperCAmelCase			:      Optional[int]			=			(
				    batch_size,
				    pipe.decoder.config.in_channels,
				    pipe.decoder.config.sample_size,
				    pipe.decoder.config.sample_size,
				)
				UpperCAmelCase			:      Any			=			pipe.prepare_latents(
				    _SCREAMING_SNAKE_CASE   ,  dtype=_SCREAMING_SNAKE_CASE   ,  device=_SCREAMING_SNAKE_CASE   ,  generator=_SCREAMING_SNAKE_CASE   ,  latents=_SCREAMING_SNAKE_CASE   ,  scheduler=DummyScheduler()		)

				UpperCAmelCase			:      Any			=			(
				    batch_size,
				    pipe.super_res_first.config.in_channels // 2,
				    pipe.super_res_first.config.sample_size,
				    pipe.super_res_first.config.sample_size,
				)
				UpperCAmelCase			:      Dict			=			pipe.prepare_latents(
				    _SCREAMING_SNAKE_CASE   ,  dtype=_SCREAMING_SNAKE_CASE   ,  device=_SCREAMING_SNAKE_CASE   ,  generator=_SCREAMING_SNAKE_CASE   ,  latents=_SCREAMING_SNAKE_CASE   ,  scheduler=DummyScheduler()		)

				UpperCAmelCase			:      List[str]			=			self.get_dummy_inputs(_SCREAMING_SNAKE_CASE   ,  pil_image=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      int			=			pipe(
				    **_SCREAMING_SNAKE_CASE   ,  decoder_latents=_SCREAMING_SNAKE_CASE   ,  super_res_latents=_SCREAMING_SNAKE_CASE		).images

				UpperCAmelCase			:      List[Any]			=			self.get_dummy_inputs(_SCREAMING_SNAKE_CASE   ,  pil_image=_SCREAMING_SNAKE_CASE		)
				# Don't pass image, instead pass embedding
				UpperCAmelCase			:      Optional[int]			=			pipeline_inputs.pop("""image"""		)
				UpperCAmelCase			:      Dict			=			pipe.image_encoder(_SCREAMING_SNAKE_CASE		).image_embeds

				UpperCAmelCase			:      Optional[int]			=			pipe(
				    **_SCREAMING_SNAKE_CASE   ,  decoder_latents=_SCREAMING_SNAKE_CASE   ,  super_res_latents=_SCREAMING_SNAKE_CASE   ,  image_embeddings=_SCREAMING_SNAKE_CASE   ,  ).images

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

			@skip_mps
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   str:





				'''simple docstring'''




				UpperCAmelCase			:      List[Any]			=			torch_device == '''cpu'''

				# Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
				UpperCAmelCase			:      Optional[int]			=			1E-2

				self._test_attention_slicing_forward_pass(
				    test_max_difference=_SCREAMING_SNAKE_CASE   ,  expected_max_diff=_SCREAMING_SNAKE_CASE		)

			@skip_mps
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   str:





				'''simple docstring'''




				UpperCAmelCase			:      List[Any]			=			torch_device == '''cpu'''
				UpperCAmelCase			:      Optional[int]			=			True
				UpperCAmelCase			:      Any			=			[
				    '''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 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[int]:





				'''simple docstring'''




				UpperCAmelCase			:      int			=			[
				    '''decoder_num_inference_steps''',
				    '''super_res_num_inference_steps''',
				]

				if torch_device == "mps":
					# TODO: MPS errors with larger batch sizes
					UpperCAmelCase			:      List[Any]			=			[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 			SCREAMING_SNAKE_CASE  (     self		)						->   Tuple:





				'''simple docstring'''




				return super().test_dict_tuple_outputs_equivalent()

			@skip_mps
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   List[Any]:





				'''simple docstring'''




				return super().test_save_load_local()




			@skip_mps
			def 			SCREAMING_SNAKE_CASE  (     self		)						->   Tuple:





				'''simple docstring'''




				return super().test_save_load_optional_components()



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

			def 			SCREAMING_SNAKE_CASE  (     self		)						->   int:





				'''simple docstring'''




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




			def 			SCREAMING_SNAKE_CASE  (     self		)						->   int:





				'''simple docstring'''




				UpperCAmelCase			:      Dict			=			load_image(
				    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"""		)
				UpperCAmelCase			:      int			=			load_numpy(
				    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
				    """/unclip/karlo_v1_alpha_cat_variation_fp16.npy"""		)

				UpperCAmelCase			:      Tuple			=			UnCLIPImageVariationPipeline.from_pretrained(
				    """kakaobrain/karlo-v1-alpha-image-variations"""   ,  torch_dtype=torch.floataa		)
				UpperCAmelCase			:      Any			=			pipeline.to(_SCREAMING_SNAKE_CASE		)
				pipeline.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

				UpperCAmelCase			:      Any			=			torch.Generator(device="""cpu"""		).manual_seed(0		)
				UpperCAmelCase			:      Any			=			pipeline(
				    _SCREAMING_SNAKE_CASE   ,  generator=_SCREAMING_SNAKE_CASE   ,  output_type="""np"""   ,  )

				UpperCAmelCase			:      Tuple			=			output.images[0]

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

				assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE   ,  15		)





"""simple docstring"""
import dataclasses
import re
import string
from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence, Tuple

import numpy as np

from . import residue_constants


A:     List[str]										=	Mapping[str, np.ndarray]
A:     Union[str, Any]										=	Mapping[str, Any]  # Is a nested dict.
A:     Any										=	0.01



@dataclasses.dataclass(frozen=UpperCAmelCase__				)
class    SCREAMING_SNAKE_CASE__						:
   __lowerCAmelCase						:     np.ndarray  # [num_res, num_atom_type, 3]

   # Amino-acid type for each residue represented as an integer between 0 and
   # 20, where 20 is 'X'.
   __lowerCAmelCase						:     np.ndarray  # [num_res]

   # Binary float mask to indicate presence of a particular atom. 1.0 if an atom
   # is present and 0.0 if not. This should be used for loss masking.
   __lowerCAmelCase						:     np.ndarray  # [num_res, num_atom_type]

   # Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
   __lowerCAmelCase						:     np.ndarray  # [num_res]

   # B-factors, or temperature factors, of each residue (in sq. angstroms units),
   # representing the displacement of the residue from its ground truth mean
   # value.
   __lowerCAmelCase						:     np.ndarray  # [num_res, num_atom_type]

   # Chain indices for multi-chain predictions
   __lowerCAmelCase						:     Optional[np.ndarray]		=		None

   # Optional remark about the protein. Included as a comment in output PDB
   # files
   __lowerCAmelCase						:     Optional[str]		=		None

   # Templates used to generate this protein (prediction-only)
   __lowerCAmelCase						:     Optional[Sequence[str]]		=		None

   # Chain corresponding to each parent
   __lowerCAmelCase						:     Optional[Sequence[int]]		=		None

def    _snake_case		(		UpperCamelCase		:	str		):
 UpperCAmelCase			:      str			=			R"""(\[[A-Z]+\]\n)"""
 UpperCAmelCase			:      List[str]			=			[tag.strip() for tag in re.split(UpperCamelCase    ,	UpperCamelCase		) if len(UpperCamelCase		) > 0]
 UpperCAmelCase			:      Iterator[Tuple[str, List[str]]]			=			zip(tags[0::2]    ,	[l.split("""\n"""		) for l in tags[1::2]]		)

 UpperCAmelCase			:      List[str]			=			["N", "CA", "C"]
 UpperCAmelCase			:      List[str]			=			None
 UpperCAmelCase			:      Tuple			=			None
 UpperCAmelCase			:      Tuple			=			None
 for g in groups:
  if "[PRIMARY]" == g[0]:
   UpperCAmelCase			:      Any			=			g[1][0].strip()
   for i in range(len(UpperCamelCase		)		):
    if seq[i] not in residue_constants.restypes:
     UpperCAmelCase			:      int			=			"""X"""  # FIXME: strings are immutable
   UpperCAmelCase			:      Optional[int]			=			np.array(
       [residue_constants.restype_order.get(UpperCamelCase    ,	residue_constants.restype_num		) for res_symbol in seq]		)
  elif "[TERTIARY]" == g[0]:
   UpperCAmelCase			:      List[List[float]]			=			[]
   for axis in range(3		):
    tertiary.append(list(map(UpperCamelCase    ,	g[1][axis].split()		)		)		)
   UpperCAmelCase			:      Optional[Any]			=			np.array(UpperCamelCase		)
   UpperCAmelCase			:      Any			=			np.zeros((len(tertiary[0]		) // 3, residue_constants.atom_type_num, 3)		).astype(np.floataa		)
   for i, atom in enumerate(UpperCamelCase		):
    UpperCAmelCase			:      Any			=			np.transpose(tertiary_np[:, i::3]		)
   atom_positions *= PICO_TO_ANGSTROM
  elif "[MASK]" == g[0]:
   UpperCAmelCase			:      Union[str, Any]			=			np.array(list(map({"""-""": 0, """+""": 1}.get    ,	g[1][0].strip()		)		)		)
   UpperCAmelCase			:      Union[str, Any]			=			np.zeros(
       (
           len(UpperCamelCase		),
           residue_constants.atom_type_num,
       )		).astype(np.floataa		)
   for i, atom in enumerate(UpperCamelCase		):
    UpperCAmelCase			:      Any			=			1
   atom_mask *= mask[..., None]

 assert aatype is not None

 return Protein(
     atom_positions=UpperCamelCase    ,	atom_mask=UpperCamelCase    ,	aatype=UpperCamelCase    ,	residue_index=np.arange(len(UpperCamelCase		)		)    ,	b_factors=UpperCamelCase    ,	)

def    _snake_case		(		UpperCamelCase		:	Protein    ,	UpperCamelCase		:	int = 0		):
 UpperCAmelCase			:      List[str]			=			[]

 UpperCAmelCase			:      List[Any]			=			prot.remark
 if remark is not None:
  pdb_headers.append(F"REMARK {remark}"		)

 UpperCAmelCase			:      List[str]			=			prot.parents
 UpperCAmelCase			:      List[Any]			=			prot.parents_chain_index
 if parents is not None and parents_chain_index is not None:
  UpperCAmelCase			:      int			=			[p for i, p in zip(UpperCamelCase    ,	UpperCamelCase		) if i == chain_id]

 if parents is None or len(UpperCamelCase		) == 0:
  UpperCAmelCase			:      Tuple			=			["""N/A"""]

 pdb_headers.append(F"PARENT {' '.join(UpperCamelCase		)}"		)

 return pdb_headers

def    _snake_case		(		UpperCamelCase		:	Protein    ,	UpperCamelCase		:	str		):
 UpperCAmelCase			:      List[str]			=			[]
 UpperCAmelCase			:      Tuple			=			pdb_str.split("""\n"""		)

 UpperCAmelCase			:      List[Any]			=			prot.remark
 if remark is not None:
  out_pdb_lines.append(F"REMARK {remark}"		)

 UpperCAmelCase			:      List[List[str]]
 if prot.parents is not None and len(prot.parents		) > 0:
  UpperCAmelCase			:      Tuple			=			[]
  if prot.parents_chain_index is not None:
   UpperCAmelCase			:      Dict[str, List[str]]			=			{}
   for p, i in zip(prot.parents    ,	prot.parents_chain_index		):
    parent_dict.setdefault(str(UpperCamelCase		)    ,	[]		)
    parent_dict[str(UpperCamelCase		)].append(UpperCamelCase		)

   UpperCAmelCase			:      Tuple			=			max([int(UpperCamelCase		) for chain_idx in parent_dict]		)
   for i in range(max_idx + 1		):
    UpperCAmelCase			:      int			=			parent_dict.get(str(UpperCamelCase		)    ,	["""N/A"""]		)
    parents_per_chain.append(UpperCamelCase		)
  else:
   parents_per_chain.append(list(prot.parents		)		)
 else:
  UpperCAmelCase			:      Union[str, Any]			=			[["""N/A"""]]

 def make_parent_line(UpperCamelCase		:	Sequence[str]		) -> str:
  return F"PARENT {' '.join(UpperCamelCase		)}"

 out_pdb_lines.append(make_parent_line(parents_per_chain[0]		)		)

 UpperCAmelCase			:      Any			=			0
 for i, l in enumerate(UpperCamelCase		):
  if "PARENT" not in l and "REMARK" not in l:
   out_pdb_lines.append(UpperCamelCase		)
  if "TER" in l and "END" not in lines[i + 1]:
   chain_counter += 1
   if not chain_counter >= len(UpperCamelCase		):
    UpperCAmelCase			:      List[Any]			=			parents_per_chain[chain_counter]
   else:
    UpperCAmelCase			:      Optional[Any]			=			["""N/A"""]

   out_pdb_lines.append(make_parent_line(UpperCamelCase		)		)

 return "\n".join(UpperCamelCase		)

def    _snake_case		(		UpperCamelCase		:	Protein		):
 UpperCAmelCase			:      int			=			residue_constants.restypes + ["""X"""]

 def res_atoa(UpperCamelCase		:	int		) -> str:
  return residue_constants.restype_atoa.get(restypes[r]    ,	"""UNK"""		)

 UpperCAmelCase			:      List[str]			=			residue_constants.atom_types

 UpperCAmelCase			:      List[str]			=			[]

 UpperCAmelCase			:      Dict			=			prot.atom_mask
 UpperCAmelCase			:      Optional[int]			=			prot.aatype
 UpperCAmelCase			:      Optional[int]			=			prot.atom_positions
 UpperCAmelCase			:      str			=			prot.residue_index.astype(np.intaa		)
 UpperCAmelCase			:      Tuple			=			prot.b_factors
 UpperCAmelCase			:      str			=			prot.chain_index

 if np.any(aatype > residue_constants.restype_num		):
  raise ValueError("""Invalid aatypes."""		)

 UpperCAmelCase			:      Any			=			get_pdb_headers(UpperCamelCase		)
 if len(UpperCamelCase		) > 0:
  pdb_lines.extend(UpperCamelCase		)

 UpperCAmelCase			:      List[str]			=			aatype.shape[0]
 UpperCAmelCase			:      int			=			1
 UpperCAmelCase			:      Optional[Any]			=			0
 UpperCAmelCase			:      str			=			string.ascii_uppercase
 UpperCAmelCase			:      Tuple			=			None
 # Add all atom sites.
 for i in range(UpperCamelCase		):
  UpperCAmelCase			:      Optional[int]			=			res_atoa(aatype[i]		)
  for atom_name, pos, mask, b_factor in zip(UpperCamelCase    ,	atom_positions[i]    ,	atom_mask[i]    ,	b_factors[i]		):
   if mask < 0.5:
    continue

   UpperCAmelCase			:      List[Any]			=			"""ATOM"""
   UpperCAmelCase			:      Dict			=			atom_name if len(UpperCamelCase		) == 4 else F" {atom_name}"
   UpperCAmelCase			:      Union[str, Any]			=			""""""
   UpperCAmelCase			:      Dict			=			""""""
   UpperCAmelCase			:      Optional[int]			=			1.00
   UpperCAmelCase			:      Union[str, Any]			=			atom_name[0]  # Protein supports only C, N, O, S, this works.
   UpperCAmelCase			:      Any			=			""""""

   UpperCAmelCase			:      List[Any]			=			"""A"""
   if chain_index is not None:
    UpperCAmelCase			:      List[str]			=			chain_tags[chain_index[i]]

   # PDB is a columnar format, every space matters here!
   UpperCAmelCase			:      Dict			=			(
       F"{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}"
       F"{res_name_a:>3} {chain_tag:>1}"
       F"{residue_index[i]:>4}{insertion_code:>1}   "
       F"{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}"
       F"{occupancy:>6.2f}{b_factor:>6.2f}          "
       F"{element:>2}{charge:>2}"
   )
   pdb_lines.append(UpperCamelCase		)
   atom_index += 1

  UpperCAmelCase			:      Union[str, Any]			=			i == n - 1
  if chain_index is not None:
   if i != n - 1 and chain_index[i + 1] != prev_chain_index:
    UpperCAmelCase			:      Any			=			True
    UpperCAmelCase			:      List[str]			=			chain_index[i + 1]

  if should_terminate:
   # Close the chain.
   UpperCAmelCase			:      Optional[int]			=			"""TER"""
   UpperCAmelCase			:      Union[str, Any]			=			(
       F"{chain_end:<6}{atom_index:>5}      {res_atoa(aatype[i]		):>3} {chain_tag:>1}{residue_index[i]:>4}"
   )
   pdb_lines.append(UpperCamelCase		)
   atom_index += 1

   if i != n - 1:
    # "prev" is a misnomer here. This happens at the beginning of
    # each new chain.
    pdb_lines.extend(get_pdb_headers(UpperCamelCase    ,	UpperCamelCase		)		)

 pdb_lines.append("""END"""		)
 pdb_lines.append(""""""		)
 return "\n".join(UpperCamelCase		)

def    _snake_case		(		UpperCamelCase		:	Protein		):
 return residue_constants.STANDARD_ATOM_MASK[prot.aatype]

def    _snake_case		(		UpperCamelCase		:	FeatureDict    ,	UpperCamelCase		:	ModelOutput    ,	UpperCamelCase		:	Optional[np.ndarray] = None    ,	UpperCamelCase		:	Optional[np.ndarray] = None    ,	UpperCamelCase		:	Optional[str] = None    ,	UpperCamelCase		:	Optional[Sequence[str]] = None    ,	UpperCamelCase		:	Optional[Sequence[int]] = None    ,	):
 return Protein(
     aatype=features["""aatype"""]    ,	atom_positions=result["""final_atom_positions"""]    ,	atom_mask=result["""final_atom_mask"""]    ,	residue_index=features["""residue_index"""] + 1    ,	b_factors=b_factors if b_factors is not None else np.zeros_like(result["""final_atom_mask"""]		)    ,	chain_index=UpperCamelCase    ,	remark=UpperCamelCase    ,	parents=UpperCamelCase    ,	parents_chain_index=UpperCamelCase    ,	)
