Fix ESMFold precompile by renaming model type

Author: murrellb

scripts/test.jl

@@ -49,7 +49,7 @@ function _load_model()
         )
 
         cfg = ESMFold.ESMFoldConfig(; trunk=trunk_cfg, lddt_head_hid_dim=lddt_head_hid_dim, use_esm_attn_map=false)
-        model = ESMFold.ESMFold(esm; cfg=cfg)
+        model = ESMFold.ESMFoldModel(esm; cfg=cfg)
         ESMFold.load_esmfold_safetensors!(model, reader)
         return model
     end

src/ESMFold.jl

@@ -20,7 +20,7 @@ export ESM2Config, ESM2
 export ESMFoldEmbedConfig, ESMFoldEmbed
 export FoldingTrunkConfig, FoldingTrunk
 export StructureModuleConfig
-export ESMFoldConfig, ESMFold
+export ESMFoldConfig, ESMFoldModel
 export load_esmfold_npz!, load_esm2_npz!, load_esmfold_safetensors!
 export load_ESM, load_ESMFold
 export sequence_to_af2_indices

src/esmfold_full.jl

@@ -35,7 +35,7 @@ function (m::ESMFoldLDDTHead)(x)
     return x
 end
 
-@concrete struct ESMFold <: Onion.Layer
+@concrete struct ESMFoldModel <: Onion.Layer
     cfg::ESMFoldConfig
     embed::ESMFoldEmbed
     trunk::FoldingTrunk
@@ -47,9 +47,9 @@ end
     lddt_head
 end
 
-@layer ESMFold
+@layer ESMFoldModel
 
-function ESMFold(
+function ESMFoldModel(
     esm::ESM2;
     cfg::ESMFoldConfig=ESMFoldConfig(),
 )
@@ -71,7 +71,7 @@ function ESMFold(
     lddt_bins = 50
     lddt_head = ESMFoldLDDTHead(cfg.trunk.structure_module.c_s, cfg.lddt_head_hid_dim, 37 * lddt_bins)
 
-    return ESMFold(
+    return ESMFoldModel(
         cfg,
         embed,
         trunk,
@@ -92,7 +92,7 @@ function _default_residx(aa::AbstractArray)
     return to_device(residx, aa, eltype(residx))
 end
 
-function (m::ESMFold)(
+function (m::ESMFoldModel)(
     aa::AbstractArray{Int,2};
     mask = nothing,
     residx = nothing,
@@ -190,7 +190,7 @@ function (m::ESMFold)(
 end
 
 function infer(
-    m::ESMFold,
+    m::ESMFoldModel,
     sequences::Union{AbstractString,AbstractVector{<:AbstractString}};
     residx = nothing,
     masking_pattern = nothing,
@@ -240,16 +240,16 @@ function infer(
     return output
 end
 
-function output_to_pdb(m::ESMFold, output::AbstractDict)
+function output_to_pdb(m::ESMFoldModel, output::AbstractDict)
     return output_to_pdb(output)
 end
 
-function infer_pdbs(m::ESMFold, seqs::AbstractVector{<:AbstractString}; kwargs...)
+function infer_pdbs(m::ESMFoldModel, seqs::AbstractVector{<:AbstractString}; kwargs...)
     output = infer(m, seqs; kwargs...)
     return output_to_pdb(output)
 end
 
-function infer_pdb(m::ESMFold, seq::AbstractString; kwargs...)
+function infer_pdb(m::ESMFoldModel, seq::AbstractString; kwargs...)
     return infer_pdbs(m, [seq]; kwargs...)[1]
 end
 
@@ -264,11 +264,11 @@ function confidence_metrics(output::AbstractDict)
     )
 end
 
-function set_chunk_size!(m::ESMFold, chunk_size::Union{Nothing,Int})
+function set_chunk_size!(m::ESMFoldModel, chunk_size::Union{Nothing,Int})
     set_chunk_size!(m.trunk, chunk_size)
     return m
 end
 
-function device_ref(m::ESMFold)
+function device_ref(m::ESMFoldModel)
     return m.embed.esm.embed_tokens.weight
 end

src/weights.jl

@@ -225,7 +225,7 @@ function _load_structure_module!(reader::SafeTensors.Reader, prefix::String, sm:
     end
 end
 
-function load_esmfold_safetensors!(model::ESMFold, reader::SafeTensors.Reader)
+function load_esmfold_safetensors!(model::ESMFoldModel, reader::SafeTensors.Reader)
     load_esmfold_safetensors!(model.embed, reader)
 
     _load_embedding_weight!(reader, "trunk.pairwise_positional_embedding.embedding.weight", model.trunk.pairwise_positional_embedding.embedding)
@@ -276,7 +276,7 @@ function load_esmfold_safetensors!(model::ESMFold, reader::SafeTensors.Reader)
     return model
 end
 
-function load_esmfold_safetensors!(model::ESMFold, path::AbstractString)
+function load_esmfold_safetensors!(model::ESMFoldModel, path::AbstractString)
     reader = SafeTensors.Reader(path)
     return load_esmfold_safetensors!(model, reader)
 end
@@ -458,7 +458,7 @@ function load_ESMFold(;
     )
     cfg = ESMFoldConfig(; trunk=trunk_cfg, lddt_head_hid_dim=lddt_head_hid_dim, use_esm_attn_map=use_esm_attn_map)
 
-    model = ESMFold(esm; cfg=cfg)
+    model = ESMFoldModel(esm; cfg=cfg)
     load_esmfold_safetensors!(model, reader)
     return model
 end