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, ModelOutput};
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 esm::models::esmc::{ESMC, ESMCConfig};
14pub use esm2::esm2::{ESM2, ESM2Config};
15pub use esm2::esm2_runner::{ESM2Models, ESM2Runner};
16pub use featurize::StructureFeatures;
17pub use ligandmpnn::configs::ProteinMPNNConfig;
18pub use ligandmpnn::model::ProteinMPNN;
19
20pub mod amplify;
21pub mod esm;
22pub mod esm2;
23pub mod featurize;
24pub mod ligandmpnn;
25pub mod types;
26
27/// Returns the best available device for computation.
28///
29/// Prioritizes CUDA GPU if available, then Metal GPU on supported platforms,
30/// and falls back to CPU if no GPU acceleration is available.
31pub fn device() -> Result<Device> {
32    if cuda_is_available() {
33        Ok(Device::new_cuda(0)?)
34    } else if metal_is_available() {
35        Ok(Device::new_metal(0)?)
36    } else {
37        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
38        {
39            println!(
40                "Running on CPU, to run on GPU(metal), build this example with `--features metal`"
41            );
42        }
43        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
44        {
45            println!("Running on CPU, to run on GPU, build this example with `--features cuda`");
46        }
47        Ok(Device::Cpu)
48    }
49}
ferritin_plms/lib.rs

ferritin_plms/
lib.rs
