koharu-diffusion 0.62.0-alpha.1

Safe Rust bindings for stable-diffusion.cpp with runtime dynamic loading.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

use std::{ffi::CStr, path::Path, ptr};

use crate::{Result, ffi::NativeCall, ffi::path_cstring, sys};

/// A ggml backend device accepted by backend assignment specifications.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Device {
    pub name: String,
    pub description: String,
}

fn copy_native_string(pointer: *const std::os::raw::c_char) -> String {
    if pointer.is_null() {
        String::new()
    } else {
        unsafe { CStr::from_ptr(pointer) }
            .to_string_lossy()
            .into_owned()
    }
}

/// Number of physical CPU cores estimated by stable-diffusion.cpp.
#[must_use]
pub fn physical_core_count() -> i32 {
    let _call = NativeCall::enter();
    unsafe { sys::sd_get_num_physical_cores() }
}

/// Native ggml/system information string.
#[must_use]
pub fn system_info() -> String {
    let _call = NativeCall::enter();
    copy_native_string(unsafe { sys::sd_get_system_info() })
}

/// stable-diffusion.cpp version string.
#[must_use]
pub fn version() -> String {
    let _call = NativeCall::enter();
    copy_native_string(unsafe { sys::sd_version() })
}

/// Source commit reported by stable-diffusion.cpp.
#[must_use]
pub fn commit() -> String {
    let _call = NativeCall::enter();
    copy_native_string(unsafe { sys::sd_commit() })
}

/// Asks ggml to discover and register every dynamic backend in `path`.
pub fn load_all_backends_from_path(path: impl AsRef<Path>) -> Result<()> {
    let path = path_cstring(path.as_ref(), "backend path")?;
    let _call = NativeCall::enter();
    unsafe { sys::ggml_backend_load_all_from_path(path.as_ptr()) };
    Ok(())
}

/// Lists backend devices as name/description pairs.
#[must_use]
pub fn list_devices() -> Vec<Device> {
    let _call = NativeCall::enter();
    let mut required = unsafe { sys::sd_list_devices(ptr::null_mut(), 0) };
    if required == 0 {
        return Vec::new();
    }

    let bytes = loop {
        let Some(buffer_len) = required.checked_add(1) else {
            return Vec::new();
        };
        let mut buffer = vec![0_u8; buffer_len];
        let reported = unsafe { sys::sd_list_devices(buffer.as_mut_ptr().cast(), buffer.len()) };
        if reported <= required {
            buffer.truncate(reported.min(buffer.len().saturating_sub(1)));
            break buffer;
        }
        required = reported;
    };

    String::from_utf8_lossy(&bytes)
        .lines()
        .map(|line| {
            let (name, description) = line.split_once('\t').unwrap_or((line, ""));
            Device {
                name: name.to_owned(),
                description: description.to_owned(),
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    #[test]
    #[ignore = "requires a stable-diffusion.cpp dynamic library on the loader path"]
    fn native_metadata_smoke_test() {
        use crate::{RngType, clear_log_callback, set_log_callback};

        assert!(!super::version().is_empty());
        assert!(super::physical_core_count() > 0);
        let _ = super::system_info();
        let _ = super::list_devices();
        assert!(matches!(RngType::parse_native("cpu"), Ok(RngType::Cpu)));
        assert_eq!(RngType::Cpu.native_name(), "cpu");
        set_log_callback(|_| {}).expect("log callback should install");
        clear_log_callback().expect("log callback should clear");
    }
}