code: 
  
  
  
  
  
  
  """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()
  A_			  =    logging.get_logger(__name__)
  
  A_			  =    [
      (
          "protein1",
          "MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA",
      ),
      ("protein2", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA"),
      ("protein3", "MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG"),
      ("protein4", "MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA"),
  ]
  
  A_			  =    {
      "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,
  }
  
  A_			  =    list("ARNDCQEGHILKMFPSTWYV")
  
  A_			  =    restypes + ["X"]
  A_			  =    restypes_with_x + ["<pad>", "<mask>", "<cls>", "<sep>", "<eos>"]
  
  
  
  
  def 						_UpperCamelCase			(  ):
  
  
  
  
  
  
  
     with TemporaryDirectory() as tempdir:
        UpperCamelCase_     ="\n".join(A			)
        UpperCamelCase_     =Path(A			) / "vocab.txt"
        vocab_file.write_text(A			)
        UpperCamelCase_     =EsmTokenizer(vocab_file=str(A			)			)
     UpperCamelCase_     =0  # Overlaps with 'A' but that seems to be what they want
     return hf_tokenizer
  
  
  
  
  def 						_UpperCamelCase			(  A    ,       A			):
     UpperCamelCase_     =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			(  A    ,       A    ,       A    ,       A    ,       A			):
  
  
  
  
  
  
  
     if model.startswith("esmfold"			):
        UpperCamelCase_     =MODEL_MAPPING[model]()
     else:
        UpperCamelCase_    ,							UpperCamelCase_     =MODEL_MAPPING[model]()
     esm.eval()  # disable dropout
  
     if model.startswith("esmfold"			):
        UpperCamelCase_     =esm.esm.embed_dim
        UpperCamelCase_     =esm.esm.num_layers
        UpperCamelCase_     =esm.esm.attention_heads
        UpperCamelCase_     =4 * embed_dim
        UpperCamelCase_     =esm.esm.token_dropout
        UpperCamelCase_     =False  # This code path does not exist in ESM-2
        UpperCamelCase_     ="rotary"
        UpperCamelCase_     =True
        UpperCamelCase_     =EsmFoldConfig()
        for key, val in esm.cfg.items():
           if hasattr(A    ,       A			) and key != "trunk":
              setattr(A    ,       A    ,       A			)
        for key, val in esm.cfg.trunk.items():
           if hasattr(esmfold_config.trunk    ,       A			) and key != "structure_module":
              setattr(esmfold_config.trunk    ,       A    ,       A			)
        for key, val in esm.cfg.trunk.structure_module.items():
           if hasattr(esmfold_config.trunk.structure_module    ,       A			):
              setattr(esmfold_config.trunk.structure_module    ,       A    ,       A			)
     elif hasattr(A    ,       "args"			):
        # Indicates an ESM-1b or ESM-1v model
        UpperCamelCase_     =esm.args.embed_dim
        UpperCamelCase_     =esm.args.layers
        UpperCamelCase_     =esm.args.attention_heads
        UpperCamelCase_     =esm.args.ffn_embed_dim
        UpperCamelCase_     =esm.args.token_dropout
        UpperCamelCase_     =True if esm.emb_layer_norm_before else False
        UpperCamelCase_     ="absolute"
        UpperCamelCase_     =False
        UpperCamelCase_     =None
     else:
        # Indicates an ESM-2 model
        UpperCamelCase_     =esm.embed_dim
        UpperCamelCase_     =esm.num_layers
        UpperCamelCase_     =esm.attention_heads
        UpperCamelCase_     =4 * embed_dim  # This is hardcoded in ESM-2
        UpperCamelCase_     =esm.token_dropout
        UpperCamelCase_     =False  # This code path does not exist in ESM-2
        UpperCamelCase_     ="rotary"
        UpperCamelCase_     =False
        UpperCamelCase_     =None
  
     if is_folding_model:
        UpperCamelCase_     =esm.esm.alphabet
     UpperCamelCase_     =tuple(alphabet.all_toks			)
     UpperCamelCase_     =alphabet.mask_idx
     UpperCamelCase_     =alphabet.padding_idx
  
     if is_folding_model:
        UpperCamelCase_     =esm.esm
     else:
        UpperCamelCase_     =esm
  
     UpperCamelCase_     =EsmConfig(
         vocab_size=original_esm_model.embed_tokens.num_embeddings    ,       mask_token_id=A    ,       hidden_size=A    ,       num_hidden_layers=A    ,       num_attention_heads=A    ,       intermediate_size=A    ,       max_position_embeddings=1_026    ,       layer_norm_eps=1e-5    ,       attention_probs_dropout_prob=0.0    ,       hidden_dropout_prob=0.0    ,       pad_token_id=A    ,       emb_layer_norm_before=A    ,       token_dropout=A    ,       position_embedding_type=A    ,       is_folding_model=A    ,       esmfold_config=A    ,       vocab_list=A    ,       )
     if classification_head:
        UpperCamelCase_     =esm.classification_heads["mnli"].out_proj.weight.shape[0]
     print("Our ESM config:"    ,       A			)
  
     if model.startswith("esmfold"			):
        UpperCamelCase_     =EsmForProteinFolding
     elif classification_head:
        UpperCamelCase_     =EsmForSequenceClassification
     else:
        UpperCamelCase_     =EsmForMaskedLM
     UpperCamelCase_     =model_class(A			)
     model.eval()
  
     # Now let's copy all the weights.
     # Embeddings
     UpperCamelCase_     =original_esm_model.embed_tokens.weight
     if position_embedding_type == "absolute":
        UpperCamelCase_     =original_esm_model.embed_positions.weight
  
     if config.emb_layer_norm_before:
        UpperCamelCase_     =original_esm_model.emb_layer_norm_before.weight
        UpperCamelCase_     =original_esm_model.emb_layer_norm_before.bias
  
     UpperCamelCase_     =original_esm_model.emb_layer_norm_after.weight
     UpperCamelCase_     =original_esm_model.emb_layer_norm_after.bias
  
     for i in range(config.num_hidden_layers			):
        # Encoder: start of layer
        UpperCamelCase_     =model.esm.encoder.layer[i]
        # esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
        UpperCamelCase_     =original_esm_model.layers[i]
  
        # self attention
        UpperCamelCase_     =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)			)
        )
  
        UpperCamelCase_     =esm_layer.self_attn.q_proj.weight
        UpperCamelCase_     =esm_layer.self_attn.q_proj.bias
        UpperCamelCase_     =esm_layer.self_attn.k_proj.weight
        UpperCamelCase_     =esm_layer.self_attn.k_proj.bias
        UpperCamelCase_     =esm_layer.self_attn.v_proj.weight
        UpperCamelCase_     =esm_layer.self_attn.v_proj.bias
  
        if getattr(esm_layer.self_attn    ,       "rot_emb"    ,       A			) 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.
           UpperCamelCase_     =esm_layer.self_attn.rot_emb.inv_freq
  
        # LayerNorm changes for pre-activation
        UpperCamelCase_     =esm_layer.self_attn_layer_norm.weight
        UpperCamelCase_     =esm_layer.self_attn_layer_norm.bias
        UpperCamelCase_     =esm_layer.final_layer_norm.weight
        UpperCamelCase_     =esm_layer.final_layer_norm.bias
  
        # self-attention output
        UpperCamelCase_     =layer.attention.output
        assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
        UpperCamelCase_     =esm_layer.self_attn.out_proj.weight
        UpperCamelCase_     =esm_layer.self_attn.out_proj.bias
  
        # intermediate
        UpperCamelCase_     =layer.intermediate
        assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
        UpperCamelCase_     =esm_layer.fca.weight
        UpperCamelCase_     =esm_layer.fca.bias
  
        # output
        UpperCamelCase_     =layer.output
        assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
        UpperCamelCase_     =esm_layer.fca.weight
        UpperCamelCase_     =esm_layer.fca.bias
        # end of layer
  
     if is_folding_model:
        UpperCamelCase_     =esm.esm_s_combine.data
        UpperCamelCase_     =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:
        UpperCamelCase_     =esm.esm.classification_heads["mnli"].dense.weight
        UpperCamelCase_     =esm.classification_heads["mnli"].dense.bias
        UpperCamelCase_     =esm.classification_heads["mnli"].out_proj.weight
        UpperCamelCase_     =esm.classification_heads["mnli"].out_proj.bias
     else:
        # LM Head
        UpperCamelCase_     =esm.lm_head.dense.weight
        UpperCamelCase_     =esm.lm_head.dense.bias
        UpperCamelCase_     =esm.lm_head.layer_norm.weight
        UpperCamelCase_     =esm.lm_head.layer_norm.bias
        UpperCamelCase_     =esm.lm_head.weight
        UpperCamelCase_     =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.
        UpperCamelCase_     =SAMPLE_DATA[:2]
     else:
        UpperCamelCase_     =SAMPLE_DATA
  
     if is_folding_model:
        UpperCamelCase_     =get_esmfold_tokenizer()
        UpperCamelCase_     =hf_tokenizer(
            [row[1] for row in sample_data]    ,       return_tensors="pt"    ,       padding=A    ,       add_special_tokens=A			)
        UpperCamelCase_    ,							UpperCamelCase_    ,							UpperCamelCase_    ,							UpperCamelCase_    ,							UpperCamelCase_     =esmfold_encode_sequences([row[1] for row in sample_data]			)
        UpperCamelCase_     =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.
        UpperCamelCase_     =alphabet.get_batch_converter()
        UpperCamelCase_    ,							UpperCamelCase_    ,							UpperCamelCase_     =batch_converter(A			)
        # Prepare tokenizer and make sure it matches
        with TemporaryDirectory() as tempdir:
           UpperCamelCase_     ="\n".join(alphabet.all_toks			)
           UpperCamelCase_     =Path(A			) / "vocab.txt"
           vocab_file.write_text(A			)
           UpperCamelCase_     =EsmTokenizer(vocab_file=str(A			)			)
  
        UpperCamelCase_     =hf_tokenizer([row[1] for row in sample_data]    ,       return_tensors="pt"    ,       padding=A			)
        UpperCamelCase_     =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.
           UpperCamelCase_     =esm.cuda().infer([row[1] for row in sample_data]			)
           UpperCamelCase_     =model.cuda()(
               input_ids=hf_tokens["input_ids"].cuda()    ,       attention_mask=hf_tokens["attention_mask"].cuda()			)
        else:
           UpperCamelCase_     =model(**A    ,       output_hidden_states=A			)
           UpperCamelCase_     =our_output["logits"]
           if classification_head:
              UpperCamelCase_     =esm.model.classification_heads["mnli"](esm.extract_features(A			)			)
           else:
              UpperCamelCase_     =esm(hf_tokens["input_ids"]    ,       repr_layers=list(range(999			)			)			)
              UpperCamelCase_     =their_output["logits"]
  
        if is_folding_model:
           UpperCamelCase_     =torch.max(torch.abs(our_output["positions"] - their_output["positions"]			)			).item()
           UpperCamelCase_     =torch.allclose(our_output["positions"]    ,       their_output["positions"]    ,       atol=1e-5			)
        else:
           UpperCamelCase_     =torch.max(torch.abs(our_output - their_output			)			).item()
           UpperCamelCase_     =torch.allclose(A    ,       A    ,       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
           UpperCamelCase_     =model.predict_contacts(hf_tokens["input_ids"]    ,       hf_tokens["attention_mask"]			)
           UpperCamelCase_     =esm.predict_contacts(hf_tokens["input_ids"]			)
           UpperCamelCase_     =torch.max(torch.abs(our_output - their_output			)			).item()
           UpperCamelCase_     =torch.allclose(A    ,       A    ,       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(A			).mkdir(parents=A    ,       exist_ok=A			)
        print(f"""Saving model to {pytorch_dump_folder_path}"""			)
        model.save_pretrained(A			)
  
        del esm  # Free up some memory before continuing
  
     print(f"""Saving tokenizer to {pytorch_dump_folder_path}"""			)
     hf_tokenizer.save_pretrained(A			)
  
     if push_to_repo:
        model.push_to_hub(repo_id=A    ,       use_auth_token=A			)
        hf_tokenizer.push_to_hub(repo_id=A    ,       use_auth_token=A			)
  
  
  if __name__ == "__main__":
      A_			  =    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.")
      A_			  =    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
      )
  
style_context: 
  
  
  
  
  
  
  """simple docstring"""
  
  
  
  
  
  
  import argparse
  import os
  
  import torch
  from torchvision.datasets.utils import download_url
  
  from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, TransformeraDModel
  
  
  A_			  =    {512: "DiT-XL-2-512x512.pt", 256: "DiT-XL-2-256x256.pt"}
  
  
  
  
  def 						_UpperCamelCase			(  A			):
     UpperCamelCase_     =f"""pretrained_models/{model_name}"""
     if not os.path.isfile(A			):
        os.makedirs("pretrained_models"    ,       exist_ok=A			)
        UpperCamelCase_     =f"""https://dl.fbaipublicfiles.com/DiT/models/{model_name}"""
        download_url(A    ,       "pretrained_models"			)
     UpperCamelCase_     =torch.load(A    ,       map_location=lambda A    ,       A			: storage			)
     return model
  
  
  
  
  def 						_UpperCamelCase			(  A			):
     UpperCamelCase_     =download_model(pretrained_models[args.image_size]			)
  
     UpperCamelCase_     =state_dict["x_embedder.proj.weight"]
     UpperCamelCase_     =state_dict["x_embedder.proj.bias"]
     state_dict.pop("x_embedder.proj.weight"			)
     state_dict.pop("x_embedder.proj.bias"			)
  
     for depth in range(28			):
        UpperCamelCase_     =state_dict[
            "t_embedder.mlp.0.weight"
        ]
        UpperCamelCase_     =state_dict[
            "t_embedder.mlp.0.bias"
        ]
        UpperCamelCase_     =state_dict[
            "t_embedder.mlp.2.weight"
        ]
        UpperCamelCase_     =state_dict[
            "t_embedder.mlp.2.bias"
        ]
        UpperCamelCase_     =state_dict[
            "y_embedder.embedding_table.weight"
        ]
  
        UpperCamelCase_     =state_dict[
            f"""blocks.{depth}.adaLN_modulation.1.weight"""
        ]
        UpperCamelCase_     =state_dict[
            f"""blocks.{depth}.adaLN_modulation.1.bias"""
        ]
  
        UpperCamelCase_    ,							UpperCamelCase_    ,							UpperCamelCase_     =torch.chunk(state_dict[f"""blocks.{depth}.attn.qkv.weight"""]    ,       3    ,       dim=0			)
        UpperCamelCase_    ,							UpperCamelCase_    ,							UpperCamelCase_     =torch.chunk(state_dict[f"""blocks.{depth}.attn.qkv.bias"""]    ,       3    ,       dim=0			)
  
        UpperCamelCase_     =q
        UpperCamelCase_     =q_bias
        UpperCamelCase_     =k
        UpperCamelCase_     =k_bias
        UpperCamelCase_     =v
        UpperCamelCase_     =v_bias
  
        UpperCamelCase_     =state_dict[
            f"""blocks.{depth}.attn.proj.weight"""
        ]
        UpperCamelCase_     =state_dict[f"""blocks.{depth}.attn.proj.bias"""]
  
        UpperCamelCase_     =state_dict[f"""blocks.{depth}.mlp.fc1.weight"""]
        UpperCamelCase_     =state_dict[f"""blocks.{depth}.mlp.fc1.bias"""]
        UpperCamelCase_     =state_dict[f"""blocks.{depth}.mlp.fc2.weight"""]
        UpperCamelCase_     =state_dict[f"""blocks.{depth}.mlp.fc2.bias"""]
  
        state_dict.pop(f"""blocks.{depth}.attn.qkv.weight"""			)
        state_dict.pop(f"""blocks.{depth}.attn.qkv.bias"""			)
        state_dict.pop(f"""blocks.{depth}.attn.proj.weight"""			)
        state_dict.pop(f"""blocks.{depth}.attn.proj.bias"""			)
        state_dict.pop(f"""blocks.{depth}.mlp.fc1.weight"""			)
        state_dict.pop(f"""blocks.{depth}.mlp.fc1.bias"""			)
        state_dict.pop(f"""blocks.{depth}.mlp.fc2.weight"""			)
        state_dict.pop(f"""blocks.{depth}.mlp.fc2.bias"""			)
        state_dict.pop(f"""blocks.{depth}.adaLN_modulation.1.weight"""			)
        state_dict.pop(f"""blocks.{depth}.adaLN_modulation.1.bias"""			)
  
     state_dict.pop("t_embedder.mlp.0.weight"			)
     state_dict.pop("t_embedder.mlp.0.bias"			)
     state_dict.pop("t_embedder.mlp.2.weight"			)
     state_dict.pop("t_embedder.mlp.2.bias"			)
     state_dict.pop("y_embedder.embedding_table.weight"			)
  
     UpperCamelCase_     =state_dict["final_layer.adaLN_modulation.1.weight"]
     UpperCamelCase_     =state_dict["final_layer.adaLN_modulation.1.bias"]
     UpperCamelCase_     =state_dict["final_layer.linear.weight"]
     UpperCamelCase_     =state_dict["final_layer.linear.bias"]
  
     state_dict.pop("final_layer.linear.weight"			)
     state_dict.pop("final_layer.linear.bias"			)
     state_dict.pop("final_layer.adaLN_modulation.1.weight"			)
     state_dict.pop("final_layer.adaLN_modulation.1.bias"			)
  
     # DiT XL/2
     UpperCamelCase_     =TransformeraDModel(
         sample_size=args.image_size // 8    ,       num_layers=28    ,       attention_head_dim=72    ,       in_channels=4    ,       out_channels=8    ,       patch_size=2    ,       attention_bias=A    ,       num_attention_heads=16    ,       activation_fn="gelu-approximate"    ,       num_embeds_ada_norm=1_000    ,       norm_type="ada_norm_zero"    ,       norm_elementwise_affine=A    ,       )
     transformer.load_state_dict(A    ,       strict=A			)
  
     UpperCamelCase_     =DDIMScheduler(
         num_train_timesteps=1_000    ,       beta_schedule="linear"    ,       prediction_type="epsilon"    ,       clip_sample=A    ,       )
  
     UpperCamelCase_     =AutoencoderKL.from_pretrained(args.vae_model			)
  
     UpperCamelCase_     =DiTPipeline(transformer=A    ,       vae=A    ,       scheduler=A			)
  
     if args.save:
        pipeline.save_pretrained(args.checkpoint_path			)
  
  
  if __name__ == "__main__":
      A_			  =    argparse.ArgumentParser()
  
      parser.add_argument(
          "--image_size",
          default=256,
          type=int,
          required=False,
          help="Image size of pretrained model, either 256 or 512.",
      )
      parser.add_argument(
          "--vae_model",
          default="stabilityai/sd-vae-ft-ema",
          type=str,
          required=False,
          help="Path to pretrained VAE model, either stabilityai/sd-vae-ft-mse or stabilityai/sd-vae-ft-ema.",
      )
      parser.add_argument(
          "--save", default=True, type=bool, required=False, help="Whether to save the converted pipeline or not."
      )
      parser.add_argument(
          "--checkpoint_path", default=None, type=str, required=True, help="Path to the output pipeline."
      )
  
      A_			  =    parser.parse_args()
      main(args)
  