






"""simple docstring"""





import argparse
import pathlib
from pathlib import Path
from tempfile import TemporaryDirectory

import esm as esm_module
import torch
from esm.esmfold.va.misc import batch_encode_sequences as esmfold_encode_sequences
from esm.esmfold.va.pretrained import esmfold_va

from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
from transformers.models.esm.modeling_esm import (
    EsmForMaskedLM,
    EsmForSequenceClassification,
    EsmIntermediate,
    EsmLayer,
    EsmOutput,
    EsmSelfAttention,
    EsmSelfOutput,
)
from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
from transformers.models.esm.tokenization_esm import EsmTokenizer
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase =			logging.get_logger(__name__)

UpperCamelCase =			[
    (
        '''protein1''',
        '''MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA''',
    ),
    ('''protein2''', '''MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA'''),
    ('''protein3''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG'''),
    ('''protein4''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA'''),
]

UpperCamelCase =			{
    '''esm1b_t33_650M_UR50S''': esm_module.pretrained.esm1b_t33_650M_UR50S,
    '''esm1v_t33_650M_UR90S_1''': esm_module.pretrained.esm1v_t33_650M_UR90S_1,
    '''esm1v_t33_650M_UR90S_2''': esm_module.pretrained.esm1v_t33_650M_UR90S_2,
    '''esm1v_t33_650M_UR90S_3''': esm_module.pretrained.esm1v_t33_650M_UR90S_3,
    '''esm1v_t33_650M_UR90S_4''': esm_module.pretrained.esm1v_t33_650M_UR90S_4,
    '''esm1v_t33_650M_UR90S_5''': esm_module.pretrained.esm1v_t33_650M_UR90S_5,
    '''esm2_t48_15B_UR50D''': esm_module.pretrained.esm2_t48_15B_UR50D,
    '''esm2_t36_3B_UR50D''': esm_module.pretrained.esm2_t36_3B_UR50D,
    '''esm2_t33_650M_UR50D''': esm_module.pretrained.esm2_t33_650M_UR50D,
    '''esm2_t30_150M_UR50D''': esm_module.pretrained.esm2_t30_150M_UR50D,
    '''esm2_t12_35M_UR50D''': esm_module.pretrained.esm2_t12_35M_UR50D,
    '''esm2_t6_8M_UR50D''': esm_module.pretrained.esm2_t6_8M_UR50D,
    '''esmfold_v1''': esmfold_va,
}

UpperCamelCase =			list("""ARNDCQEGHILKMFPSTWYV""")

UpperCamelCase =			restypes + ['''X''']
UpperCamelCase =			restypes_with_x + ['''<pad>''', '''<mask>''', '''<cls>''', '''<sep>''', '''<eos>''']






def 				A(  ):
       """simple docstring"""





       with TemporaryDirectory() as tempdir:
              lowercase__:				List[Any]			  =			"""\n""".join(lowerCAmelCase_		)
              lowercase__:				Optional[int]			  =			Path(lowerCAmelCase_		) / """vocab.txt"""
              vocab_file.write_text(lowerCAmelCase_		)
              lowercase__:				Optional[int]			  =			EsmTokenizer(vocab_file=str(lowerCAmelCase_		)		)
       lowercase__:				Union[str, Any]			  =			0  # Overlaps with 'A' but that seems to be what they want
       return hf_tokenizer






def 				A(  snake_case_						,   snake_case_		):
       """simple docstring"""





       lowercase__:				Dict			  =			our_module.load_state_dict(original_module.state_dict()		)
       if status.missing_keys:
              raise ValueError(F"""Missing keys: {status.missing_keys}"""		)
       if status.unexpected_keys:
              raise ValueError(F"""Unexpected keys: {status.unexpected_keys}"""		)






def 				A(  snake_case_						,   snake_case_						,   snake_case_						,   snake_case_						,   snake_case_		):
       """simple docstring"""





       if model.startswith("esmfold"		):
              lowercase__:				List[Any]			  =			MODEL_MAPPING[model]()
       else:
              lowercase__:				int			  =			MODEL_MAPPING[model]()
       esm.eval()  # disable dropout

       if model.startswith("esmfold"		):
              lowercase__:				Optional[Any]			  =			esm.esm.embed_dim
              lowercase__:				Optional[Any]			  =			esm.esm.num_layers
              lowercase__:				List[Any]			  =			esm.esm.attention_heads
              lowercase__:				Any			  =			4 * embed_dim
              lowercase__:				int			  =			esm.esm.token_dropout
              lowercase__:				List[str]			  =			False  # This code path does not exist in ESM-2
              lowercase__:				int			  =			"""rotary"""
              lowercase__:				str			  =			True
              lowercase__:				Any			  =			EsmFoldConfig()
              for key, val in esm.cfg.items():
                     if hasattr(lowerCAmelCase_						,   lowerCAmelCase_		) and key != "trunk":
                            setattr(lowerCAmelCase_						,   lowerCAmelCase_						,   lowerCAmelCase_		)
              for key, val in esm.cfg.trunk.items():
                     if hasattr(esmfold_config.trunk						,   lowerCAmelCase_		) and key != "structure_module":
                            setattr(esmfold_config.trunk						,   lowerCAmelCase_						,   lowerCAmelCase_		)
              for key, val in esm.cfg.trunk.structure_module.items():
                     if hasattr(esmfold_config.trunk.structure_module						,   lowerCAmelCase_		):
                            setattr(esmfold_config.trunk.structure_module						,   lowerCAmelCase_						,   lowerCAmelCase_		)
       elif hasattr(lowerCAmelCase_						,   "args"		):
              # Indicates an ESM-1b or ESM-1v model
              lowercase__:				Any			  =			esm.args.embed_dim
              lowercase__:				List[str]			  =			esm.args.layers
              lowercase__:				str			  =			esm.args.attention_heads
              lowercase__:				List[Any]			  =			esm.args.ffn_embed_dim
              lowercase__:				List[Any]			  =			esm.args.token_dropout
              lowercase__:				Dict			  =			True if esm.emb_layer_norm_before else False
              lowercase__:				Optional[Any]			  =			"""absolute"""
              lowercase__:				List[Any]			  =			False
              lowercase__:				Optional[int]			  =			None
       else:
              # Indicates an ESM-2 model
              lowercase__:				Dict			  =			esm.embed_dim
              lowercase__:				Tuple			  =			esm.num_layers
              lowercase__:				int			  =			esm.attention_heads
              lowercase__:				List[Any]			  =			4 * embed_dim  # This is hardcoded in ESM-2
              lowercase__:				Dict			  =			esm.token_dropout
              lowercase__:				Union[str, Any]			  =			False  # This code path does not exist in ESM-2
              lowercase__:				List[str]			  =			"""rotary"""
              lowercase__:				Tuple			  =			False
              lowercase__:				Any			  =			None

       if is_folding_model:
              lowercase__:				int			  =			esm.esm.alphabet
       lowercase__:				int			  =			tuple(alphabet.all_toks		)
       lowercase__:				Dict			  =			alphabet.mask_idx
       lowercase__:				Dict			  =			alphabet.padding_idx

       if is_folding_model:
              lowercase__:				Dict			  =			esm.esm
       else:
              lowercase__:				Any			  =			esm

       lowercase__:				str			  =			EsmConfig(
           vocab_size=original_esm_model.embed_tokens.num_embeddings						,   mask_token_id=lowerCAmelCase_						,   hidden_size=lowerCAmelCase_						,   num_hidden_layers=lowerCAmelCase_						,   num_attention_heads=lowerCAmelCase_						,   intermediate_size=lowerCAmelCase_						,   max_position_embeddings=1026						,   layer_norm_eps=1E-5						,   attention_probs_dropout_prob=0.0						,   hidden_dropout_prob=0.0						,   pad_token_id=lowerCAmelCase_						,   emb_layer_norm_before=lowerCAmelCase_						,   token_dropout=lowerCAmelCase_						,   position_embedding_type=lowerCAmelCase_						,   is_folding_model=lowerCAmelCase_						,   esmfold_config=lowerCAmelCase_						,   vocab_list=lowerCAmelCase_						,   )
       if classification_head:
              lowercase__:				Optional[Any]			  =			esm.classification_heads["""mnli"""].out_proj.weight.shape[0]
       print("Our ESM config:"						,   lowerCAmelCase_		)

       if model.startswith("esmfold"		):
              lowercase__:				int			  =			EsmForProteinFolding
       elif classification_head:
              lowercase__:				Optional[int]			  =			EsmForSequenceClassification
       else:
              lowercase__:				Dict			  =			EsmForMaskedLM
       lowercase__:				List[Any]			  =			model_class(lowerCAmelCase_		)
       model.eval()

       # Now let's copy all the weights.
       # Embeddings
       lowercase__:				List[Any]			  =			original_esm_model.embed_tokens.weight
       if position_embedding_type == "absolute":
              lowercase__:				List[str]			  =			original_esm_model.embed_positions.weight

       if config.emb_layer_norm_before:
              lowercase__:				int			  =			original_esm_model.emb_layer_norm_before.weight
              lowercase__:				Tuple			  =			original_esm_model.emb_layer_norm_before.bias

       lowercase__:				Optional[int]			  =			original_esm_model.emb_layer_norm_after.weight
       lowercase__:				List[Any]			  =			original_esm_model.emb_layer_norm_after.bias

       for i in range(config.num_hidden_layers		):
              # Encoder: start of layer
              lowercase__:				EsmLayer			  =			model.esm.encoder.layer[i]
              # esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
              lowercase__:				int			  =			original_esm_model.layers[i]

              # self attention
              lowercase__:				EsmSelfAttention			  =			layer.attention.self
              assert (
                  esm_layer.self_attn.k_proj.weight.data.shape
                  == esm_layer.self_attn.q_proj.weight.data.shape
                  == esm_layer.self_attn.v_proj.weight.data.shape
                  == torch.Size((config.hidden_size, config.hidden_size)		)
              )

              lowercase__:				Optional[int]			  =			esm_layer.self_attn.q_proj.weight
              lowercase__:				Dict			  =			esm_layer.self_attn.q_proj.bias
              lowercase__:				Optional[Any]			  =			esm_layer.self_attn.k_proj.weight
              lowercase__:				Union[str, Any]			  =			esm_layer.self_attn.k_proj.bias
              lowercase__:				List[Any]			  =			esm_layer.self_attn.v_proj.weight
              lowercase__:				Optional[Any]			  =			esm_layer.self_attn.v_proj.bias

              if getattr(esm_layer.self_attn						,   "rot_emb"						,   lowerCAmelCase_		) is not None:
                     # Matt: Although inv_freq is not a trainable weight, it is computed at model init and cached.
                     # During the training of ESM-2 the model was converted to float16 precision, which also converts
                     # the inv_freq tensor, and the loss of precision remains even if the model is loaded later as float32.
                     # If we recompute inv_freq without this loss of precision then we will get subtly different rotary
                     # embeddings, which are enough to cause significant discrepancies in model outputs. To avoid this,
                     # we make sure the new model copies the data from the old inv_freq.
                     lowercase__:				Any			  =			esm_layer.self_attn.rot_emb.inv_freq

              # LayerNorm changes for pre-activation
              lowercase__:				Dict			  =			esm_layer.self_attn_layer_norm.weight
              lowercase__:				List[Any]			  =			esm_layer.self_attn_layer_norm.bias
              lowercase__:				Tuple			  =			esm_layer.final_layer_norm.weight
              lowercase__:				Union[str, Any]			  =			esm_layer.final_layer_norm.bias

              # self-attention output
              lowercase__:				EsmSelfOutput			  =			layer.attention.output
              assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
              lowercase__:				Optional[int]			  =			esm_layer.self_attn.out_proj.weight
              lowercase__:				str			  =			esm_layer.self_attn.out_proj.bias

              # intermediate
              lowercase__:				EsmIntermediate			  =			layer.intermediate
              assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
              lowercase__:				int			  =			esm_layer.fca.weight
              lowercase__:				Tuple			  =			esm_layer.fca.bias

              # output
              lowercase__:				EsmOutput			  =			layer.output
              assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
              lowercase__:				Tuple			  =			esm_layer.fca.weight
              lowercase__:				Optional[int]			  =			esm_layer.fca.bias
              # end of layer

       if is_folding_model:
              lowercase__:				Optional[int]			  =			esm.esm_s_combine.data
              lowercase__:				int			  =			esm.afa_to_esm.data
              transfer_and_check_weights(esm.embedding						,   model.embedding		)
              transfer_and_check_weights(esm.esm_s_mlp						,   model.esm_s_mlp		)
              transfer_and_check_weights(esm.trunk						,   model.trunk		)
              transfer_and_check_weights(esm.distogram_head						,   model.distogram_head		)
              transfer_and_check_weights(esm.ptm_head						,   model.ptm_head		)
              transfer_and_check_weights(esm.lm_head						,   model.lm_head		)
              transfer_and_check_weights(esm.lddt_head						,   model.lddt_head		)

       elif classification_head:
              lowercase__:				Tuple			  =			esm.esm.classification_heads["""mnli"""].dense.weight
              lowercase__:				Dict			  =			esm.classification_heads["""mnli"""].dense.bias
              lowercase__:				List[str]			  =			esm.classification_heads["""mnli"""].out_proj.weight
              lowercase__:				List[str]			  =			esm.classification_heads["""mnli"""].out_proj.bias
       else:
              # LM Head
              lowercase__:				Union[str, Any]			  =			esm.lm_head.dense.weight
              lowercase__:				Optional[Any]			  =			esm.lm_head.dense.bias
              lowercase__:				Optional[Any]			  =			esm.lm_head.layer_norm.weight
              lowercase__:				Tuple			  =			esm.lm_head.layer_norm.bias
              lowercase__:				Tuple			  =			esm.lm_head.weight
              lowercase__:				Any			  =			esm.lm_head.bias

       # Contact prediction head
       transfer_and_check_weights(esm.contact_head						,   model.esm.contact_head		)

       # Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
       if is_folding_model:
              # Folding models aren't trained on masked inputs and don't like mask tokens.
              lowercase__:				Optional[Any]			  =			SAMPLE_DATA[:2]
       else:
              lowercase__:				Optional[Any]			  =			SAMPLE_DATA

       if is_folding_model:
              lowercase__:				int			  =			get_esmfold_tokenizer()
              lowercase__:				Optional[int]			  =			hf_tokenizer(
                  [row[1] for row in sample_data]						,   return_tensors="pt"						,   padding=lowerCAmelCase_						,   add_special_tokens=lowerCAmelCase_		)
              lowercase__:				Dict			  =			esmfold_encode_sequences([row[1] for row in sample_data]		)
              lowercase__:				Union[str, Any]			  =			torch.all(hf_tokens["input_ids"] == esmfold_aas		) and torch.all(
                  hf_tokens["attention_mask"] == esmfold_mask		)
       else:
              # Let's check that we get the same results.
              lowercase__:				Optional[Any]			  =			alphabet.get_batch_converter()
              lowercase__:				Optional[int]			  =			batch_converter(lowerCAmelCase_		)
              # Prepare tokenizer and make sure it matches
              with TemporaryDirectory() as tempdir:
                     lowercase__:				Union[str, Any]			  =			"""\n""".join(alphabet.all_toks		)
                     lowercase__:				List[str]			  =			Path(lowerCAmelCase_		) / """vocab.txt"""
                     vocab_file.write_text(lowerCAmelCase_		)
                     lowercase__:				Any			  =			EsmTokenizer(vocab_file=str(lowerCAmelCase_		)		)

              lowercase__:				Union[str, Any]			  =			hf_tokenizer([row[1] for row in sample_data]						,   return_tensors="pt"						,   padding=lowerCAmelCase_		)
              lowercase__:				Dict			  =			torch.all(hf_tokens["input_ids"] == batch_tokens		)

       print("Do both models tokenizers output the same tokens?"						,   "🔥" if success else "💩"		)
       if not success:
              raise Exception("Tokenization does not match!"		)

       with torch.no_grad():
              if is_folding_model:
                     # Let's test the model in parts
                     # ESMFold always converts the ESM stem to float16, which requires float16 ops
                     # that don't exist on CPU. Therefore, to test it we need to run it on GPU. However,
                     # ESMFold is what we in the community call a "big boy" and so we desperately avoid putting both the
                     # original and the converted model on the GPU at the same time.
                     lowercase__:				Union[str, Any]			  =			esm.cuda().infer([row[1] for row in sample_data]		)
                     lowercase__:				List[str]			  =			model.cuda()(
                         input_ids=hf_tokens["input_ids"].cuda()						,   attention_mask=hf_tokens["attention_mask"].cuda()		)
              else:
                     lowercase__:				Any			  =			model(**lowerCAmelCase_						,   output_hidden_states=lowerCAmelCase_		)
                     lowercase__:				List[str]			  =			our_output["""logits"""]
                     if classification_head:
                            lowercase__:				Tuple			  =			esm.model.classification_heads["""mnli"""](esm.extract_features(lowerCAmelCase_		)		)
                     else:
                            lowercase__:				Optional[Any]			  =			esm(hf_tokens["input_ids"]						,   repr_layers=list(range(999		)		)		)
                            lowercase__:				int			  =			their_output["""logits"""]

              if is_folding_model:
                     lowercase__:				Any			  =			torch.max(torch.abs(our_output["positions"] - their_output["positions"]		)		).item()
                     lowercase__:				str			  =			torch.allclose(our_output["positions"]						,   their_output["positions"]						,   atol=1E-5		)
              else:
                     lowercase__:				List[Any]			  =			torch.max(torch.abs(our_output - their_output		)		).item()
                     lowercase__:				List[Any]			  =			torch.allclose(lowerCAmelCase_						,   lowerCAmelCase_						,   atol=1E-5		)

              print(F"""max_absolute_diff = {max_absolute_diff}"""		)  # ~ 1e-5
              print("Do both models output the same tensors?"						,   "🔥" if success else "💩"		)

              if not success:
                     raise Exception("Something went wRoNg"		)

              if not is_folding_model:
                     # Let's check contact prediction too
                     lowercase__:				Union[str, Any]			  =			model.predict_contacts(hf_tokens["input_ids"]						,   hf_tokens["attention_mask"]		)
                     lowercase__:				Tuple			  =			esm.predict_contacts(hf_tokens["input_ids"]		)
                     lowercase__:				Optional[int]			  =			torch.max(torch.abs(our_output - their_output		)		).item()
                     lowercase__:				Optional[int]			  =			torch.allclose(lowerCAmelCase_						,   lowerCAmelCase_						,   atol=1E-5		)

                     print("Contact prediction testing:"		)
                     print(F"""max_absolute_diff = {max_absolute_diff}"""		)  # ~ 1e-5
                     print("Do both models output the same tensors?"						,   "🔥" if success else "💩"		)

                     if not success:
                            raise Exception("Something went wRoNg"		)

              pathlib.Path(lowerCAmelCase_		).mkdir(parents=lowerCAmelCase_						,   exist_ok=lowerCAmelCase_		)
              print(F"""Saving model to {pytorch_dump_folder_path}"""		)
              model.save_pretrained(lowerCAmelCase_		)

              del esm  # Free up some memory before continuing

       print(F"""Saving tokenizer to {pytorch_dump_folder_path}"""		)
       hf_tokenizer.save_pretrained(lowerCAmelCase_		)

       if push_to_repo:
              model.push_to_hub(repo_id=lowerCAmelCase_						,   use_auth_token=lowerCAmelCase_		)
              hf_tokenizer.push_to_hub(repo_id=lowerCAmelCase_						,   use_auth_token=lowerCAmelCase_		)


if __name__ == "__main__":
  UpperCamelCase =			argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--pytorch_dump_folder_path""", type=str, required=True, help="""Path to the output PyTorch model."""
  )
  parser.add_argument(
      """--classification_head""", action="""store_true""", help="""Whether to convert a final classification head."""
  )
  parser.add_argument("""--model""", default=None, type=str, required=True, help="""Name of model to convert.""")
  parser.add_argument("""--push_to_repo""", type=str, help="""Repo to upload to (including username!).""")
  parser.add_argument("""--auth_token""", type=str, help="""HuggingFace auth token.""")
  UpperCamelCase =			parser.parse_args()
  convert_esm_checkpoint_to_pytorch(
      args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
  )








'''simple docstring'''

import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest







class 	__lowerCAmelCase    (  __a				):
     snake_case   :    List[str]							      =      (HeunDiscreteScheduler,)
     snake_case   :    Optional[Any]							      =      1_0



     def     snake_case_ (self						,	**lowerCAmelCase__	):
          _UpperCAmelCase					:  Union[str, Any]         =    {
              """num_train_timesteps""": 1_1_0_0,
              """beta_start""": 0.0_0_0_1,
              """beta_end""": 0.0_2,
              """beta_schedule""": """linear""",
          }

          config.update(**lowerCAmelCase__	)
          return config



     def     snake_case_ (self	):
          for timesteps in [1_0, 5_0, 1_0_0, 1_0_0_0]:
               self.check_over_configs(num_train_timesteps=lowerCAmelCase__	)



     def     snake_case_ (self	):
          for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1]						,	[0.0_0_0_2, 0.0_0_2, 0.0_2]	):
               self.check_over_configs(beta_start=lowerCAmelCase__						,	beta_end=lowerCAmelCase__	)



     def     snake_case_ (self	):
          for schedule in ["linear", "scaled_linear", "exp"]:
               self.check_over_configs(beta_schedule=lowerCAmelCase__	)



     def     snake_case_ (self	):
          for clip_sample_range in [1.0, 2.0, 3.0]:
               self.check_over_configs(clip_sample_range=lowerCAmelCase__						,	clip_sample=lowerCAmelCase__	)



     def     snake_case_ (self	):
          for prediction_type in ["epsilon", "v_prediction", "sample"]:
               self.check_over_configs(prediction_type=lowerCAmelCase__	)



     def     snake_case_ (self	):
          _UpperCAmelCase					:  Tuple         =    self.scheduler_classes[0]
          _UpperCAmelCase					:  Optional[int]         =    self.get_scheduler_config()
          _UpperCAmelCase					:  List[str]         =    scheduler_class(**lowerCAmelCase__	)

          scheduler.set_timesteps(self.num_inference_steps	)

          _UpperCAmelCase					:  str         =    self.dummy_model()
          _UpperCAmelCase					:  List[Any]         =    self.dummy_sample_deter * scheduler.init_noise_sigma
          _UpperCAmelCase					:  Tuple         =    sample.to(lowerCAmelCase__	)

          for i, t in enumerate(scheduler.timesteps	):
               _UpperCAmelCase					:  Union[str, Any]         =    scheduler.scale_model_input(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Any         =    model(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Tuple         =    scheduler.step(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	)
               _UpperCAmelCase					:  Optional[Any]         =    output.prev_sample

          _UpperCAmelCase					:  Optional[Any]         =    torch.sum(torch.abs(lowerCAmelCase__	)	)
          _UpperCAmelCase					:  int         =    torch.mean(torch.abs(lowerCAmelCase__	)	)

          if torch_device in ["cpu", "mps"]:
               assert abs(result_sum.item() - 0.1_2_3_3	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3
          else:
               # CUDA
               assert abs(result_sum.item() - 0.1_2_3_3	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3



     def     snake_case_ (self	):
          _UpperCAmelCase					:  Optional[int]         =    self.scheduler_classes[0]
          _UpperCAmelCase					:  Any         =    self.get_scheduler_config(prediction_type="""v_prediction"""	)
          _UpperCAmelCase					:  Any         =    scheduler_class(**lowerCAmelCase__	)

          scheduler.set_timesteps(self.num_inference_steps	)

          _UpperCAmelCase					:  Optional[Any]         =    self.dummy_model()
          _UpperCAmelCase					:  Any         =    self.dummy_sample_deter * scheduler.init_noise_sigma
          _UpperCAmelCase					:  Any         =    sample.to(lowerCAmelCase__	)

          for i, t in enumerate(scheduler.timesteps	):
               _UpperCAmelCase					:  List[str]         =    scheduler.scale_model_input(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  List[Any]         =    model(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[Any]         =    scheduler.step(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	)
               _UpperCAmelCase					:  Union[str, Any]         =    output.prev_sample

          _UpperCAmelCase					:  Union[str, Any]         =    torch.sum(torch.abs(lowerCAmelCase__	)	)
          _UpperCAmelCase					:  str         =    torch.mean(torch.abs(lowerCAmelCase__	)	)

          if torch_device in ["cpu", "mps"]:
               assert abs(result_sum.item() - 4.6_934e-07	) < 1e-2
               assert abs(result_mean.item() - 6.1_112e-10	) < 1e-3
          else:
               # CUDA
               assert abs(result_sum.item() - 4.693_428_650_170_972e-07	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3



     def     snake_case_ (self	):
          _UpperCAmelCase					:  Tuple         =    self.scheduler_classes[0]
          _UpperCAmelCase					:  Any         =    self.get_scheduler_config()
          _UpperCAmelCase					:  Optional[int]         =    scheduler_class(**lowerCAmelCase__	)

          scheduler.set_timesteps(self.num_inference_steps						,	device=lowerCAmelCase__	)

          _UpperCAmelCase					:  List[Any]         =    self.dummy_model()
          _UpperCAmelCase					:  Dict         =    self.dummy_sample_deter.to(lowerCAmelCase__	) * scheduler.init_noise_sigma

          for t in scheduler.timesteps:
               _UpperCAmelCase					:  Union[str, Any]         =    scheduler.scale_model_input(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[Any]         =    model(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[int]         =    scheduler.step(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	)
               _UpperCAmelCase					:  int         =    output.prev_sample

          _UpperCAmelCase					:  Any         =    torch.sum(torch.abs(lowerCAmelCase__	)	)
          _UpperCAmelCase					:  Optional[Any]         =    torch.mean(torch.abs(lowerCAmelCase__	)	)

          if str(lowerCAmelCase__	).startswith("""cpu"""	):
               # The following sum varies between 148 and 156 on mps. Why?
               assert abs(result_sum.item() - 0.1_2_3_3	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3
          elif str(lowerCAmelCase__	).startswith("""mps"""	):
               # Larger tolerance on mps
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-2
          else:
               # CUDA
               assert abs(result_sum.item() - 0.1_2_3_3	) < 1e-2
               assert abs(result_mean.item() - 0.0_0_0_2	) < 1e-3



     def     snake_case_ (self	):
          _UpperCAmelCase					:  int         =    self.scheduler_classes[0]
          _UpperCAmelCase					:  str         =    self.get_scheduler_config()
          _UpperCAmelCase					:  Union[str, Any]         =    scheduler_class(**lowerCAmelCase__						,	use_karras_sigmas=lowerCAmelCase__	)

          scheduler.set_timesteps(self.num_inference_steps						,	device=lowerCAmelCase__	)

          _UpperCAmelCase					:  Optional[int]         =    self.dummy_model()
          _UpperCAmelCase					:  Optional[Any]         =    self.dummy_sample_deter.to(lowerCAmelCase__	) * scheduler.init_noise_sigma
          _UpperCAmelCase					:  Union[str, Any]         =    sample.to(lowerCAmelCase__	)

          for t in scheduler.timesteps:
               _UpperCAmelCase					:  Dict         =    scheduler.scale_model_input(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[Any]         =    model(lowerCAmelCase__						,	lowerCAmelCase__	)

               _UpperCAmelCase					:  Optional[int]         =    scheduler.step(lowerCAmelCase__						,	lowerCAmelCase__						,	lowerCAmelCase__	)
               _UpperCAmelCase					:  Optional[Any]         =    output.prev_sample

          _UpperCAmelCase					:  Union[str, Any]         =    torch.sum(torch.abs(lowerCAmelCase__	)	)
          _UpperCAmelCase					:  List[Any]         =    torch.mean(torch.abs(lowerCAmelCase__	)	)

          assert abs(result_sum.item() - 0.0_0_0_1_5	) < 1e-2
          assert abs(result_mean.item() - 1.9_869_554_535_034_695e-07	) < 1e-2


