ferritin_plms/
lib.rs

1//! ferritin-plms
2//!
3//!
4//! ```shell
5//! cargo run --example amplify
6//! cargo run --example amplify --features metal
7//! ```
8pub use amplify::amplify::{AMPLIFY, AmplifyOutput};
9pub use amplify::amplify_runner::{AmplifyModels, AmplifyRunner};
10pub use amplify::config::AMPLIFYConfig;
11use candle_core::utils::{cuda_is_available, metal_is_available};
12use candle_core::{Device, Result};
13pub use esm2::esm2::{ESM2, ESM2Config};
14pub use esm2::esm2_runner::{ESM2Models, ESM2Runner};
15pub use esm3::models::esm3::ESM3Config;
16pub use esm3::pretrained::{ESM3Models, ESM3Runner};
17pub use esmc::models::esmc::{ESMC, ESMCConfig, ESMCOutput, LogitsConfig, LogitsOutput};
18pub use esmc::pretrained::{ESMCModels, ESMCRunner};
19pub use esmfold2::config::ESMFold2Config;
20pub use esmfold2::output::ESMFold2Output;
21pub use esmfold2::pretrained::{ESMFold2Models, ESMFold2Runner};
22pub use esmfold2::{
23    ChainInput, DNAInput, LigandInput, Modification, ProteinInput, StructurePredictionInput,
24};
25pub use featurize::StructureFeatures;
26pub use ligandmpnn::configs::ProteinMPNNConfig;
27pub use ligandmpnn::model::ProteinMPNN;
28pub use ligandmpnn::pmpnn_runner::{ProteinMPNNModels, ProteinMPNNRunner};
29
30pub mod amplify;
31pub mod esm2;
32pub mod esm3;
33pub mod esmc;
34pub mod esmfold2;
35pub mod featurize;
36pub mod ligandmpnn;
37pub mod plm_runner;
38pub mod types;
39pub mod utils;
40pub use plm_runner::PlmRunner;
41
42/// Returns the best available device for computation.
43///
44/// If `cpu` is true, always returns `Device::Cpu` regardless of available hardware.
45/// Otherwise prioritizes CUDA GPU if available, then Metal GPU on supported platforms,
46/// and falls back to CPU if no GPU acceleration is available.
47pub fn device(cpu: bool) -> Result<Device> {
48    if cpu {
49        return Ok(Device::Cpu);
50    }
51    if cuda_is_available() {
52        Ok(Device::new_cuda(0)?)
53    } else if metal_is_available() {
54        Ok(Device::new_metal(0)?)
55    } else {
56        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
57        {
58            println!(
59                "Running on CPU, to run on GPU(metal), build this example with `--features metal`"
60            );
61        }
62        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
63        {
64            println!("Running on CPU, to run on GPU, build this example with `--features cuda`");
65        }
66        Ok(Device::Cpu)
67    }
68}
ferritin_plms/lib.rs

ferritin_plms/
lib.rs
