refactor: simplify device name logic

koharu-ml/src/lib.rs

@@ -8,42 +8,12 @@ pub mod llm;
 pub mod manga_ocr;
 
 use anyhow::Result;
-use candle_core::{Device, utils::metal_is_available};
+use candle_core::utils::metal_is_available;
 
+pub use candle_core::Device;
 pub use koharu_http::hf_hub::set_cache_dir;
 pub use llm::{language_from_tag, set_default_locale, set_locale, supported_locales};
 
-/// Name of the compute device being used.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub enum DeviceName {
-    Cpu,
-    Cuda,
-    Metal,
-}
-
-impl std::fmt::Display for DeviceName {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            DeviceName::Cpu => write!(f, "CPU"),
-            DeviceName::Cuda => write!(f, "CUDA"),
-            DeviceName::Metal => write!(f, "Metal"),
-        }
-    }
-}
-
-/// Returns the name of the device that would be selected.
-pub fn device_name(cpu: bool) -> DeviceName {
-    if cpu {
-        DeviceName::Cpu
-    } else if cuda_is_available() {
-        DeviceName::Cuda
-    } else if metal_is_available() {
-        DeviceName::Metal
-    } else {
-        DeviceName::Cpu
-    }
-}
-
 pub fn device(cpu: bool) -> Result<Device> {
     if cpu {
         Ok(Device::Cpu)

koharu-pipeline/src/lib.rs

@@ -3,7 +3,7 @@ pub mod pipeline;
 
 use std::sync::Arc;
 
-use koharu_ml::DeviceName;
+use koharu_ml::Device;
 use koharu_renderer::facade::Renderer;
 use koharu_types::AppState;
 use tokio::sync::RwLock;
@@ -14,7 +14,7 @@ pub struct AppResources {
     pub ml: Arc<koharu_ml::facade::Model>,
     pub llm: Arc<koharu_ml::llm::facade::Model>,
     pub renderer: Arc<Renderer>,
-    pub ml_device: DeviceName,
+    pub device: Device,
     pub pipeline: Arc<RwLock<Option<pipeline::PipelineHandle>>>,
     pub version: &'static str,
 }

koharu-pipeline/src/operations.rs

@@ -28,7 +28,11 @@ pub struct DeviceInfo {
 
 pub async fn device(state: AppResources) -> anyhow::Result<DeviceInfo> {
     Ok(DeviceInfo {
-        ml_device: state.ml_device.to_string(),
+        ml_device: match state.device {
+            koharu_ml::Device::Cpu => "CPU".to_string(),
+            koharu_ml::Device::Cuda(_) => "CUDA".to_string(),
+            koharu_ml::Device::Metal(_) => "Metal".to_string(),
+        },
         wgpu: state.renderer.wgpu_device_info(),
     })
 }

koharu/src/app.rs

@@ -8,7 +8,7 @@ use tauri::Manager;
 use tokio::{net::TcpListener, sync::RwLock};
 use tracing_subscriber::fmt::format::FmtSpan;
 
-use koharu_ml::{cuda_is_available, device_name};
+use koharu_ml::{cuda_is_available, device};
 use koharu_pipeline::AppResources;
 use koharu_renderer::facade::Renderer;
 use koharu_rpc::{SharedResources, server};
@@ -131,7 +131,6 @@ async fn build_resources(cpu: bool) -> Result<AppResources> {
         );
     }
 
-    let ml_device = device_name(cpu);
     let ml = Arc::new(
         koharu_ml::facade::Model::new(cpu)
             .await
@@ -146,7 +145,7 @@ async fn build_resources(cpu: bool) -> Result<AppResources> {
         ml,
         llm,
         renderer,
-        ml_device,
+        device: device(cpu)?,
         pipeline: Arc::new(RwLock::new(None)),
         version: crate::version::current(),
     })