_target_: models.simplefold.torch.architecture.FoldingDiT

hidden_size: 1280
num_heads: 20
atom_num_heads: 6
output_channels: 3
use_atom_mask: False
use_length_condition: True
esm_dropout_prob: 0.0
esm_model: esm2_3B

time_embedder:
  _target_: models.simplefold.torch.layers.TimestepEmbedder
  hidden_size: 1280
aminoacid_pos_embedder:
  _target_: models.simplefold.torch.pos_embed.AbsolutePositionEncoding
  in_dim: 1
  embed_dim: 1280
  include_input: True
pos_embedder:
  _target_: models.simplefold.torch.pos_embed.FourierPositionEncoding
  in_dim: 3
  include_input: True
  min_freq_log2: 0
  max_freq_log2: 12
  num_freqs: 128
  log_sampling: True

trunk:
  _target_: models.simplefold.torch.blocks.HomogenTrunk
  depth: 36
  block: 
    _target_: models.simplefold.torch.blocks.DiTBlock
    _partial_: True # because in the for loop we create a new module
    hidden_size: 1280
    mlp_ratio: 4.0
    use_swiglu: True # SwiGLU FFN
    self_attention_layer:
      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
      _partial_: True
      hidden_size: 1280
      num_heads: 20
      qk_norm: True
      pos_embedder:
        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
        in_dim: 4
        embed_dim: 1280
        num_heads: 20
        base: 100.0

atom_hidden_size_enc: 384
atom_n_queries_enc: 32
atom_n_keys_enc: 128
atom_encoder_transformer:
  _target_: models.simplefold.torch.blocks.HomogenTrunk
  depth: 2
  block: 
    _target_: models.simplefold.torch.blocks.DiTBlock
    _partial_: True # because in the for loop we create a new module
    hidden_size: 384
    mlp_ratio: 4.0
    use_swiglu: True # SwiGLU FFN
    self_attention_layer:
      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
      _partial_: True
      hidden_size: 384
      num_heads: 6
      qk_norm: True
      pos_embedder:
        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
        in_dim: 4
        embed_dim: 384
        num_heads: 6
        base: 100.0

atom_hidden_size_dec: 384
atom_n_queries_dec: 32
atom_n_keys_dec: 128
atom_decoder_transformer:
  _target_: models.simplefold.torch.blocks.HomogenTrunk
  depth: 2
  block: 
    _target_: models.simplefold.torch.blocks.DiTBlock
    _partial_: True # because in the for loop we create a new module
    hidden_size: 384
    mlp_ratio: 4.0
    use_swiglu: True # SwiGLU FFN
    self_attention_layer:
      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
      _partial_: True
      hidden_size: 384
      num_heads: 6
      qk_norm: True
      pos_embedder:
        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
        in_dim: 4
        embed_dim: 384
        num_heads: 6
        base: 100.0