use crate::engine::download;
use anyhow::{anyhow, bail, Context, Result};
use std::path::{Path, PathBuf};
use tracing::{info, warn};
const TRACE_TARGET: &str = "studio_worker::engine::sd_provision";
const DEFAULT_RELEASE_TAG: &str = "master-669-2d40a8b";
const RELEASE_ENV: &str = "STUDIO_WORKER_SDCPP_RELEASE";
const URL_ENV: &str = "STUDIO_WORKER_SDCPP_URL";
pub fn binary_name() -> &'static str {
if cfg!(target_os = "windows") {
"sd-cli.exe"
} else {
"sd-cli"
}
}
fn library_name() -> &'static str {
if cfg!(target_os = "windows") {
"stable-diffusion.dll"
} else if cfg!(target_os = "macos") {
"libstable-diffusion.dylib"
} else {
"libstable-diffusion.so"
}
}
fn vulkan_loader_name() -> Option<&'static str> {
if cfg!(target_os = "windows") {
Some("vulkan-1.dll")
} else if cfg!(target_os = "macos") {
None
} else {
Some("libvulkan.so.1")
}
}
fn vulkan_remedy() -> &'static str {
if cfg!(target_os = "windows") {
"install/update your GPU driver (NVIDIA, AMD, or Intel) — it ships \
the Vulkan runtime (vulkan-1.dll)"
} else {
"install the Vulkan loader + a GPU driver, e.g. on Debian/Ubuntu \
`sudo apt install libvulkan1 mesa-vulkan-drivers` (plus the \
vendor driver for NVIDIA/AMD); verify with `vulkaninfo --summary`"
}
}
#[cfg_attr(coverage_nightly, coverage(off))]
fn vulkan_loader_loads() -> bool {
match vulkan_loader_name() {
None => true,
Some(name) => unsafe { libloading::Library::new(name).is_ok() },
}
}
fn vulkan_runtime_status_with(loader_loads: bool) -> Result<()> {
let Some(loader) = vulkan_loader_name() else {
return Ok(()); };
if loader_loads {
return Ok(());
}
bail!(
"Vulkan runtime not available: the loader `{loader}` could not be \
loaded, so stable-diffusion.cpp cannot run on the GPU. We cannot \
auto-provision it — {}.",
vulkan_remedy()
)
}
#[cfg_attr(coverage_nightly, coverage(off))]
pub fn vulkan_runtime_status() -> Result<()> {
vulkan_runtime_status_with(vulkan_loader_loads())
}
fn release_tag() -> String {
std::env::var(RELEASE_ENV).unwrap_or_else(|_| DEFAULT_RELEASE_TAG.to_string())
}
fn sha_from_tag(tag: &str) -> Result<&str> {
match tag.rsplit_once('-') {
Some((_, sha)) if !sha.is_empty() => Ok(sha),
_ => Err(anyhow!("release tag {tag:?} has no '-<sha>' segment")),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AssetSource {
Upstream,
SelfHosted,
}
fn asset_plan(os: &str, arch: &str) -> Result<(AssetSource, &'static str)> {
use AssetSource::*;
match (os, arch) {
("windows", "x86_64") => Ok((Upstream, "win-vulkan-x64")),
("linux", "x86_64") => Ok((Upstream, "Linux-Ubuntu-24.04-x86_64-vulkan")),
("macos", "aarch64") | ("macos", "x86_64") => Ok((Upstream, "Darwin-macOS-15.7.7-arm64")),
("linux", "aarch64") => Ok((SelfHosted, "Linux-aarch64-vulkan")),
_ => bail!(
"no prebuilt stable-diffusion.cpp binary for {os}/{arch}; \
install sd-cli manually — see docs/operations/sd-cli-install.md"
),
}
}
fn asset_name(sha: &str, suffix: &str) -> String {
format!("sd-master-{sha}-bin-{suffix}.zip")
}
fn self_hosted_tag(upstream_tag: &str) -> String {
format!("sdcpp-prebuilt-{upstream_tag}")
}
fn download_url(tag: &str, os: &str, arch: &str) -> Result<String> {
let sha = sha_from_tag(tag)?;
let (source, suffix) = asset_plan(os, arch)?;
let asset = asset_name(sha, suffix);
Ok(match source {
AssetSource::Upstream => format!(
"https://github.com/leejet/stable-diffusion.cpp/releases/download/{tag}/{asset}"
),
AssetSource::SelfHosted => format!(
"https://github.com/webbertakken/studio-worker/releases/download/{}/{asset}",
self_hosted_tag(tag)
),
})
}
fn resolve_url() -> Result<String> {
if let Ok(url) = std::env::var(URL_ENV) {
if !url.is_empty() {
return Ok(url);
}
}
download_url(&release_tag(), std::env::consts::OS, std::env::consts::ARCH)
}
pub fn library_path_env(sd_cli: &Path) -> Option<(&'static str, PathBuf)> {
if cfg!(target_os = "windows") {
return None;
}
let dir = sd_cli.parent()?;
if dir.join(library_name()).is_file() {
let var = if cfg!(target_os = "macos") {
"DYLD_LIBRARY_PATH"
} else {
"LD_LIBRARY_PATH"
};
Some((var, dir.to_path_buf()))
} else {
None
}
}
#[cfg_attr(coverage_nightly, coverage(off))]
fn extract_zip(zip_path: &Path, dest_dir: &Path) -> Result<usize> {
let file =
std::fs::File::open(zip_path).with_context(|| format!("opening {}", zip_path.display()))?;
let mut archive = zip::ZipArchive::new(file)
.with_context(|| format!("reading zip {}", zip_path.display()))?;
std::fs::create_dir_all(dest_dir)
.with_context(|| format!("creating {}", dest_dir.display()))?;
let mut written = 0usize;
for i in 0..archive.len() {
let mut entry = archive.by_index(i)?;
if entry.is_dir() {
continue;
}
let Some(file_name) = Path::new(entry.name()).file_name().map(|n| n.to_owned()) else {
warn!(
target: TRACE_TARGET,
op = "extract",
name = entry.name(),
"skipping zip entry with no file name"
);
continue;
};
let out = dest_dir.join(&file_name);
let mode = entry.unix_mode();
let mut writer =
std::fs::File::create(&out).with_context(|| format!("creating {}", out.display()))?;
std::io::copy(&mut entry, &mut writer)
.with_context(|| format!("writing {}", out.display()))?;
drop(writer);
apply_unix_mode(&out, mode)?;
written += 1;
}
Ok(written)
}
#[cfg(unix)]
fn apply_unix_mode(path: &Path, mode: Option<u32>) -> Result<()> {
use std::os::unix::fs::PermissionsExt;
if let Some(mode) = mode {
std::fs::set_permissions(path, std::fs::Permissions::from_mode(mode))
.with_context(|| format!("chmod {}", path.display()))?;
}
Ok(())
}
#[cfg(not(unix))]
fn apply_unix_mode(_path: &Path, _mode: Option<u32>) -> Result<()> {
Ok(())
}
#[cfg(unix)]
fn make_executable(path: &Path) -> Result<()> {
use std::os::unix::fs::PermissionsExt;
let mut perms = std::fs::metadata(path)
.with_context(|| format!("stat {}", path.display()))?
.permissions();
perms.set_mode(perms.mode() | 0o755);
std::fs::set_permissions(path, perms).with_context(|| format!("chmod +x {}", path.display()))
}
#[cfg(not(unix))]
fn make_executable(_path: &Path) -> Result<()> {
Ok(())
}
fn install_dir(staging: &Path, target: &Path) -> Result<usize> {
std::fs::create_dir_all(target).with_context(|| format!("creating {}", target.display()))?;
let mut moved = 0usize;
for entry in
std::fs::read_dir(staging).with_context(|| format!("reading {}", staging.display()))?
{
let entry = entry?;
if !entry.file_type()?.is_file() {
continue;
}
let from = entry.path();
let to = target.join(entry.file_name());
if to.exists() {
std::fs::remove_file(&to).with_context(|| format!("replacing {}", to.display()))?;
}
if std::fs::rename(&from, &to).is_err() {
std::fs::copy(&from, &to)
.with_context(|| format!("copying {} -> {}", from.display(), to.display()))?;
}
moved += 1;
}
Ok(moved)
}
#[cfg_attr(coverage_nightly, coverage(off))]
pub fn provision(models_root: &Path) -> Result<PathBuf> {
let target_dir = models_root.join("bin");
let binary = target_dir.join(binary_name());
if binary.is_file() {
return Ok(binary);
}
let url = resolve_url()?;
info!(
target: TRACE_TARGET,
op = "provision",
url = %url,
dest = %target_dir.display(),
"sd-cli not found; provisioning stable-diffusion.cpp"
);
std::fs::create_dir_all(models_root)
.with_context(|| format!("creating {}", models_root.display()))?;
let stamp = format!("{}-{}", std::process::id(), now_nanos());
let zip_path = models_root.join(format!(".sd-cli-{stamp}.zip"));
let staging = models_root.join(format!(".sd-cli-staging-{stamp}"));
let result = (|| -> Result<PathBuf> {
download::download_file(&url, &zip_path)
.with_context(|| format!("downloading sd-cli zip from {url}"))?;
let count = extract_zip(&zip_path, &staging)?;
let staged_binary = staging.join(binary_name());
if !staged_binary.is_file() {
bail!(
"downloaded sd-cli zip from {url} did not contain {} (extracted {count} files)",
binary_name()
);
}
install_dir(&staging, &target_dir)?;
make_executable(&binary)?;
if !binary.is_file() {
bail!("sd-cli install left no binary at {}", binary.display());
}
Ok(binary.clone())
})();
let _ = std::fs::remove_file(&zip_path);
let _ = std::fs::remove_dir_all(&staging);
match &result {
Ok(path) => info!(
target: TRACE_TARGET,
op = "provision",
path = %path.display(),
"sd-cli provisioned"
),
Err(e) => warn!(
target: TRACE_TARGET,
op = "provision",
error = %e,
"sd-cli provisioning failed"
),
}
result
}
#[cfg_attr(coverage_nightly, coverage(off))]
fn now_nanos() -> i64 {
chrono::Utc::now().timestamp_nanos_opt().unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::tempdir;
#[test]
fn sha_from_tag_takes_trailing_segment() {
assert_eq!(sha_from_tag("master-669-2d40a8b").unwrap(), "2d40a8b");
assert_eq!(sha_from_tag("master-1-abc").unwrap(), "abc");
}
#[test]
fn sha_from_tag_rejects_a_tag_without_a_sha() {
assert!(sha_from_tag("master").is_err());
assert!(sha_from_tag("trailing-").is_err());
}
#[test]
fn asset_plan_picks_vulkan_or_universal_for_supported_targets() {
use AssetSource::*;
assert_eq!(
asset_plan("windows", "x86_64").unwrap(),
(Upstream, "win-vulkan-x64")
);
assert_eq!(
asset_plan("linux", "x86_64").unwrap(),
(Upstream, "Linux-Ubuntu-24.04-x86_64-vulkan")
);
assert_eq!(
asset_plan("macos", "aarch64").unwrap(),
(Upstream, "Darwin-macOS-15.7.7-arm64")
);
}
#[test]
fn asset_plan_makes_intel_mac_and_arm_linux_first_class() {
use AssetSource::*;
assert_eq!(
asset_plan("macos", "x86_64").unwrap(),
(Upstream, "Darwin-macOS-15.7.7-arm64")
);
assert_eq!(
asset_plan("linux", "aarch64").unwrap(),
(SelfHosted, "Linux-aarch64-vulkan")
);
}
#[test]
fn asset_plan_rejects_unsupported_targets_with_guidance() {
let err = asset_plan("freebsd", "x86_64").unwrap_err().to_string();
assert!(err.contains("no prebuilt"), "got: {err}");
assert!(
err.contains("sd-cli-install.md"),
"points to the doc: {err}"
);
assert!(asset_plan("windows", "aarch64").is_err());
}
#[test]
fn asset_name_embeds_sha_and_platform() {
assert_eq!(
asset_name("2d40a8b", "win-vulkan-x64"),
"sd-master-2d40a8b-bin-win-vulkan-x64.zip"
);
assert_eq!(
asset_name("2d40a8b", "Linux-aarch64-vulkan"),
"sd-master-2d40a8b-bin-Linux-aarch64-vulkan.zip"
);
}
#[test]
fn download_url_targets_upstream_for_covered_platforms() {
let url = download_url("master-669-2d40a8b", "windows", "x86_64").unwrap();
let expected = concat!(
"https://github.com/leejet/stable-diffusion.cpp/releases/download/",
"master-669-2d40a8b/sd-master-2d40a8b-bin-win-vulkan-x64.zip"
);
assert_eq!(url, expected);
}
#[test]
fn download_url_targets_our_release_for_arm_linux() {
let url = download_url("master-669-2d40a8b", "linux", "aarch64").unwrap();
let expected = concat!(
"https://github.com/webbertakken/studio-worker/releases/download/",
"sdcpp-prebuilt-master-669-2d40a8b/",
"sd-master-2d40a8b-bin-Linux-aarch64-vulkan.zip"
);
assert_eq!(url, expected);
}
#[test]
fn download_url_uses_universal_darwin_asset_for_intel_mac() {
let arm = download_url("master-669-2d40a8b", "macos", "aarch64").unwrap();
let intel = download_url("master-669-2d40a8b", "macos", "x86_64").unwrap();
assert_eq!(arm, intel, "Intel Macs use the same universal2 asset");
assert!(intel.contains("Darwin-macOS-15.7.7-arm64"), "got: {intel}");
}
#[test]
fn install_dir_moves_files_and_overwrites() {
let staging = tempdir().unwrap();
let target = tempdir().unwrap();
std::fs::write(staging.path().join("sd-cli"), b"new-binary").unwrap();
std::fs::write(staging.path().join("libstable-diffusion.so"), b"lib").unwrap();
std::fs::write(target.path().join("sd-cli"), b"old-binary").unwrap();
let moved = install_dir(staging.path(), target.path()).unwrap();
assert_eq!(moved, 2);
assert_eq!(
std::fs::read(target.path().join("sd-cli")).unwrap(),
b"new-binary"
);
assert_eq!(
std::fs::read(target.path().join("libstable-diffusion.so")).unwrap(),
b"lib"
);
assert!(!staging.path().join("sd-cli").exists());
}
#[test]
fn extract_zip_flattens_and_defuses_zip_slip() {
let dir = tempdir().unwrap();
let zip_path = dir.path().join("test.zip");
{
let file = std::fs::File::create(&zip_path).unwrap();
let mut zw = zip::ZipWriter::new(file);
let opts: zip::write::FileOptions<()> = zip::write::FileOptions::default()
.compression_method(zip::CompressionMethod::Deflated);
zw.start_file("sd-cli", opts).unwrap();
zw.write_all(b"binary").unwrap();
zw.start_file("nested/libstable-diffusion.so", opts)
.unwrap();
zw.write_all(b"lib").unwrap();
zw.start_file("../../escape.txt", opts).unwrap();
zw.write_all(b"evil").unwrap();
zw.finish().unwrap();
}
let dest = dir.path().join("out");
let count = extract_zip(&zip_path, &dest).unwrap();
assert_eq!(count, 3);
assert_eq!(std::fs::read(dest.join("sd-cli")).unwrap(), b"binary");
assert_eq!(
std::fs::read(dest.join("libstable-diffusion.so")).unwrap(),
b"lib"
);
assert!(dest.join("escape.txt").is_file());
assert!(!dir.path().join("escape.txt").exists());
}
#[cfg(unix)]
#[test]
fn extract_zip_preserves_exec_bit() {
use std::os::unix::fs::PermissionsExt;
let dir = tempdir().unwrap();
let zip_path = dir.path().join("exec.zip");
{
let file = std::fs::File::create(&zip_path).unwrap();
let mut zw = zip::ZipWriter::new(file);
let opts: zip::write::FileOptions<()> = zip::write::FileOptions::default()
.compression_method(zip::CompressionMethod::Deflated)
.unix_permissions(0o755);
zw.start_file("sd-cli", opts).unwrap();
zw.write_all(b"#!/bin/sh\n").unwrap();
zw.finish().unwrap();
}
let dest = dir.path().join("out");
extract_zip(&zip_path, &dest).unwrap();
let mode = std::fs::metadata(dest.join("sd-cli"))
.unwrap()
.permissions()
.mode();
assert!(mode & 0o111 != 0, "exec bit must survive: {mode:o}");
}
#[cfg(unix)]
#[test]
fn library_path_env_points_loader_at_sibling_lib() {
let dir = tempdir().unwrap();
let sd_cli = dir.path().join(binary_name());
std::fs::write(&sd_cli, b"bin").unwrap();
assert!(library_path_env(&sd_cli).is_none());
std::fs::write(dir.path().join(library_name()), b"lib").unwrap();
let (var, env_dir) = library_path_env(&sd_cli).expect("sibling lib resolved");
assert!(var == "LD_LIBRARY_PATH" || var == "DYLD_LIBRARY_PATH");
assert_eq!(env_dir, dir.path());
}
#[test]
fn vulkan_status_ok_when_loader_loads() {
assert!(vulkan_runtime_status_with(true).is_ok());
}
#[test]
fn vulkan_status_errors_with_actionable_remedy_when_missing() {
let result = vulkan_runtime_status_with(false);
if cfg!(target_os = "macos") {
assert!(result.is_ok());
} else {
let err = result.unwrap_err().to_string();
assert!(err.contains("Vulkan runtime"), "got: {err}");
assert!(
err.contains("auto-provision"),
"must say we can't auto-provision it: {err}"
);
if cfg!(target_os = "windows") {
assert!(err.contains("vulkan-1.dll"), "got: {err}");
assert!(err.contains("GPU driver"), "got: {err}");
} else {
assert!(err.contains("libvulkan1"), "got: {err}");
assert!(err.contains("vulkaninfo"), "got: {err}");
}
}
}
#[cfg(target_os = "windows")]
#[test]
fn library_path_env_is_none_on_windows() {
let dir = tempdir().unwrap();
let sd_cli = dir.path().join(binary_name());
std::fs::write(&sd_cli, b"bin").unwrap();
std::fs::write(dir.path().join(library_name()), b"lib").unwrap();
assert!(library_path_env(&sd_cli).is_none());
}
}