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

ferritin_plms/
lib.rs
