


"""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()
SCREAMING_SNAKE_CASE				:     List[str]          =	logging.get_logger(__name__)

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

SCREAMING_SNAKE_CASE				:     Tuple          =	{
    '''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,
}

SCREAMING_SNAKE_CASE				:     Any          =	list('''ARNDCQEGHILKMFPSTWYV''')

SCREAMING_SNAKE_CASE				:     Any          =	restypes + ['''X''']
SCREAMING_SNAKE_CASE				:     int          =	restypes_with_x + ['''<pad>''', '''<mask>''', '''<cls>''', '''<sep>''', '''<eos>''']






def 	__UpperCAmelCase					(       )							->	int:






   """simple docstring"""


   with TemporaryDirectory() as tempdir:
      _lowerCAmelCase					     =		"\n".join(UpperCamelCase__   )
      _lowerCAmelCase					     =		Path(UpperCamelCase__   ) / "vocab.txt"
      vocab_file.write_text(UpperCamelCase__   )
      _lowerCAmelCase					     =		EsmTokenizer(vocab_file=str(UpperCamelCase__   )   )
   _lowerCAmelCase					     =		0  # Overlaps with 'A' but that seems to be what they want
   return hf_tokenizer






def 	__UpperCAmelCase					(       snake_case_  :  List[Any]   ,	snake_case_  :  int   )							->	Tuple:






   """simple docstring"""


   _lowerCAmelCase					     =		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 	__UpperCAmelCase					(       snake_case_  :  str   ,	snake_case_  :  str   ,	snake_case_  :  bool   ,	snake_case_  :  str   ,	snake_case_  :  str   )							->	Tuple:






   """simple docstring"""


   if model.startswith("""esmfold"""   ):
      _lowerCAmelCase					     =		MODEL_MAPPING[model]()
   else:
      _lowerCAmelCase					     =		MODEL_MAPPING[model]()
   esm.eval()  # disable dropout

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

   if is_folding_model:
      _lowerCAmelCase					     =		esm.esm.alphabet
   _lowerCAmelCase					     =		tuple(alphabet.all_toks   )
   _lowerCAmelCase					     =		alphabet.mask_idx
   _lowerCAmelCase					     =		alphabet.padding_idx

   if is_folding_model:
      _lowerCAmelCase					     =		esm.esm
   else:
      _lowerCAmelCase					     =		esm

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

   if model.startswith("""esmfold"""   ):
      _lowerCAmelCase					     =		EsmForProteinFolding
   elif classification_head:
      _lowerCAmelCase					     =		EsmForSequenceClassification
   else:
      _lowerCAmelCase					     =		EsmForMaskedLM
   _lowerCAmelCase					     =		model_class(UpperCamelCase__   )
   model.eval()

   # Now let's copy all the weights.
   # Embeddings
   _lowerCAmelCase					     =		original_esm_model.embed_tokens.weight
   if position_embedding_type == "absolute":
      _lowerCAmelCase					     =		original_esm_model.embed_positions.weight

   if config.emb_layer_norm_before:
      _lowerCAmelCase					     =		original_esm_model.emb_layer_norm_before.weight
      _lowerCAmelCase					     =		original_esm_model.emb_layer_norm_before.bias

   _lowerCAmelCase					     =		original_esm_model.emb_layer_norm_after.weight
   _lowerCAmelCase					     =		original_esm_model.emb_layer_norm_after.bias

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

      # self attention
      _lowerCAmelCase					     =		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)   )
      )

      _lowerCAmelCase					     =		esm_layer.self_attn.q_proj.weight
      _lowerCAmelCase					     =		esm_layer.self_attn.q_proj.bias
      _lowerCAmelCase					     =		esm_layer.self_attn.k_proj.weight
      _lowerCAmelCase					     =		esm_layer.self_attn.k_proj.bias
      _lowerCAmelCase					     =		esm_layer.self_attn.v_proj.weight
      _lowerCAmelCase					     =		esm_layer.self_attn.v_proj.bias

      if getattr(esm_layer.self_attn   ,	"""rot_emb"""   ,	UpperCamelCase__   ) 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.
         _lowerCAmelCase					     =		esm_layer.self_attn.rot_emb.inv_freq

      # LayerNorm changes for pre-activation
      _lowerCAmelCase					     =		esm_layer.self_attn_layer_norm.weight
      _lowerCAmelCase					     =		esm_layer.self_attn_layer_norm.bias
      _lowerCAmelCase					     =		esm_layer.final_layer_norm.weight
      _lowerCAmelCase					     =		esm_layer.final_layer_norm.bias

      # self-attention output
      _lowerCAmelCase					     =		layer.attention.output
      assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
      _lowerCAmelCase					     =		esm_layer.self_attn.out_proj.weight
      _lowerCAmelCase					     =		esm_layer.self_attn.out_proj.bias

      # intermediate
      _lowerCAmelCase					     =		layer.intermediate
      assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
      _lowerCAmelCase					     =		esm_layer.fca.weight
      _lowerCAmelCase					     =		esm_layer.fca.bias

      # output
      _lowerCAmelCase					     =		layer.output
      assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
      _lowerCAmelCase					     =		esm_layer.fca.weight
      _lowerCAmelCase					     =		esm_layer.fca.bias
      # end of layer

   if is_folding_model:
      _lowerCAmelCase					     =		esm.esm_s_combine.data
      _lowerCAmelCase					     =		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:
      _lowerCAmelCase					     =		esm.esm.classification_heads["mnli"].dense.weight
      _lowerCAmelCase					     =		esm.classification_heads["mnli"].dense.bias
      _lowerCAmelCase					     =		esm.classification_heads["mnli"].out_proj.weight
      _lowerCAmelCase					     =		esm.classification_heads["mnli"].out_proj.bias
   else:
      # LM Head
      _lowerCAmelCase					     =		esm.lm_head.dense.weight
      _lowerCAmelCase					     =		esm.lm_head.dense.bias
      _lowerCAmelCase					     =		esm.lm_head.layer_norm.weight
      _lowerCAmelCase					     =		esm.lm_head.layer_norm.bias
      _lowerCAmelCase					     =		esm.lm_head.weight
      _lowerCAmelCase					     =		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.
      _lowerCAmelCase					     =		SAMPLE_DATA[:2]
   else:
      _lowerCAmelCase					     =		SAMPLE_DATA

   if is_folding_model:
      _lowerCAmelCase					     =		get_esmfold_tokenizer()
      _lowerCAmelCase					     =		hf_tokenizer(
          [row[1] for row in sample_data]   ,	return_tensors="""pt"""   ,	padding=UpperCamelCase__   ,	add_special_tokens=UpperCamelCase__   )
      _lowerCAmelCase					     =		esmfold_encode_sequences([row[1] for row in sample_data]   )
      _lowerCAmelCase					     =		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.
      _lowerCAmelCase					     =		alphabet.get_batch_converter()
      _lowerCAmelCase					     =		batch_converter(UpperCamelCase__   )
      # Prepare tokenizer and make sure it matches
      with TemporaryDirectory() as tempdir:
         _lowerCAmelCase					     =		"\n".join(alphabet.all_toks   )
         _lowerCAmelCase					     =		Path(UpperCamelCase__   ) / "vocab.txt"
         vocab_file.write_text(UpperCamelCase__   )
         _lowerCAmelCase					     =		EsmTokenizer(vocab_file=str(UpperCamelCase__   )   )

      _lowerCAmelCase					     =		hf_tokenizer([row[1] for row in sample_data]   ,	return_tensors="""pt"""   ,	padding=UpperCamelCase__   )
      _lowerCAmelCase					     =		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.
         _lowerCAmelCase					     =		esm.cuda().infer([row[1] for row in sample_data]   )
         _lowerCAmelCase					     =		model.cuda()(
             input_ids=hf_tokens["""input_ids"""].cuda()   ,	attention_mask=hf_tokens["""attention_mask"""].cuda()   )
      else:
         _lowerCAmelCase					     =		model(**UpperCamelCase__   ,	output_hidden_states=UpperCamelCase__   )
         _lowerCAmelCase					     =		our_output["logits"]
         if classification_head:
            _lowerCAmelCase					     =		esm.model.classification_heads["mnli"](esm.extract_features(UpperCamelCase__   )   )
         else:
            _lowerCAmelCase					     =		esm(hf_tokens["""input_ids"""]   ,	repr_layers=list(range(999   )   )   )
            _lowerCAmelCase					     =		their_output["logits"]

      if is_folding_model:
         _lowerCAmelCase					     =		torch.max(torch.abs(our_output["""positions"""] - their_output["""positions"""]   )   ).item()
         _lowerCAmelCase					     =		torch.allclose(our_output["""positions"""]   ,	their_output["""positions"""]   ,	atol=1e-5   )
      else:
         _lowerCAmelCase					     =		torch.max(torch.abs(our_output - their_output   )   ).item()
         _lowerCAmelCase					     =		torch.allclose(UpperCamelCase__   ,	UpperCamelCase__   ,	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
         _lowerCAmelCase					     =		model.predict_contacts(hf_tokens["""input_ids"""]   ,	hf_tokens["""attention_mask"""]   )
         _lowerCAmelCase					     =		esm.predict_contacts(hf_tokens["""input_ids"""]   )
         _lowerCAmelCase					     =		torch.max(torch.abs(our_output - their_output   )   ).item()
         _lowerCAmelCase					     =		torch.allclose(UpperCamelCase__   ,	UpperCamelCase__   ,	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(UpperCamelCase__   ).mkdir(parents=UpperCamelCase__   ,	exist_ok=UpperCamelCase__   )
      print(F"""Saving model to {pytorch_dump_folder_path}"""   )
      model.save_pretrained(UpperCamelCase__   )

      del esm  # Free up some memory before continuing

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

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


if __name__ == "__main__":
  SCREAMING_SNAKE_CASE				:     Union[str, Any]          =	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.''')
  SCREAMING_SNAKE_CASE				:     List[Any]          =	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	(   __lowercase    ):
   __UpperCamelCase									=   (HeunDiscreteScheduler,)
   __UpperCamelCase									=   10







   def 							A__						(self       ,	**lowerCamelCase				):




      '''simple docstring'''




      _lowerCAmelCase					     =		{
          """num_train_timesteps""": 1_100,
          """beta_start""": 0.0001,
          """beta_end""": 0.02,
          """beta_schedule""": """linear""",
      }

      config.update(**lowerCamelCase				)
      return config







   def 							A__						(self				):




      '''simple docstring'''




      for timesteps in [10, 50, 100, 1_000]:
         self.check_over_configs(num_train_timesteps=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]       ,	[0.0002, 0.002, 0.02]				):
         self.check_over_configs(beta_start=lowerCamelCase       ,	beta_end=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




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







   def 							A__						(self				):




      '''simple docstring'''




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







   def 							A__						(self				):




      '''simple docstring'''




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







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps				)

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for i, t in enumerate(scheduler.timesteps				):
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

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







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config(prediction_type="""v_prediction"""				)
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps				)

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for i, t in enumerate(scheduler.timesteps				):
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if torch_device in ["cpu", "mps"]:
         assert abs(result_sum.item() - 4.6934e-07				) < 1e-2
         assert abs(result_mean.item() - 6.1112e-10				) < 1e-3
      else:
         # CUDA
         assert abs(result_sum.item() - 4.693428650170972e-07				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

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

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter.to(lowerCamelCase				) * scheduler.init_noise_sigma

      for t in scheduler.timesteps:
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		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.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3
      elif str(lowerCamelCase				).startswith("""mps"""				):
         # Larger tolerance on mps
         assert abs(result_mean.item() - 0.0002				) < 1e-2
      else:
         # CUDA
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase       ,	use_karras_sigmas=lowerCamelCase				)

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

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter.to(lowerCamelCase				) * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for t in scheduler.timesteps:
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      assert abs(result_sum.item() - 0.0_0015				) < 1e-2
      assert abs(result_mean.item() - 1.9869554535034695e-07				) < 1e-2