Function get_inference_device

cake-core/src/utils/mod.rs:144–159

Returns the best available device at `ordinal` index (in case of multiple GPUs), or CPU if `force_cpu` is true.

(force_cpu: bool, ordinal: usize)

/// Returns the best available device at `ordinal` index (in case of multiple GPUs), or CPU if `force_cpu` is true.
pub fn get_inference_device(force_cpu: bool, ordinal: usize) -> Result<Device> {
    if force_cpu {
        log::debug!("device is forced cpu");
        Ok(Device::Cpu)
    } else if cuda_is_available() {
        log::debug!("device is cuda {ordinal}");
        Ok(Device::new_cuda(ordinal)?)
    } else if metal_is_available() {
        log::debug!("device is metal {ordinal}");
        Ok(Device::new_metal(ordinal)?)
    } else {
        log::debug!("device is cpu");
        // fallback to cpu if nothing else available
        Ok(Device::Cpu)
    }
}

pub fn load_safetensors_from_model(path: &Path) -> Result<Vec<std::path::PathBuf>> {
    log::info!("loading tensors from {} ...", "model.safetensors");