diff --git a/Cargo.toml b/Cargo.toml index ad988ea..7dafda6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,14 +7,12 @@ edition = "2021" [features] wgpu-backend = ["burn-wgpu"] +default = ["wgpu-backend"] [dependencies] burn = "0.20.1" -burn-ndarray = "0.20.1" -burn-tch = "0.20.1" burn-autodiff = "0.20.1" burn-wgpu = { version = "0.20.1", optional = true } -tch = "0.22.0" serde = {version = "1.0.171", features = ["std", "derive"]} npy = "0.4.0" num-traits = "0.2.15" diff --git a/README.md b/README.md index e92c52c..f00a7e9 100644 --- a/README.md +++ b/README.md @@ -4,8 +4,6 @@ Stable-Diffusion-Burn is a Rust-based project which ports the V1 stable diffusio ## How To Use -### Step 0: Install libtorch v2.4.1 - ### Step 1: Download the Model and Set Environment Variables Start by downloading the SDv1-4 model provided on HuggingFace. @@ -16,22 +14,17 @@ wget https://huggingface.co/Gadersd/Stable-Diffusion-Burn/resolve/main/SDv1-4.mp ### Step 2: Run the Sample Binary -Invoke the sample binary provided in the rust code. By default, torch is used. The WGPU backend is unstable for SD but may work well in the future as burn-wpu is optimized. +Invoke the sample binary provided in the rust code. The application now uses a pure Rust backend (WGPU/Vulkan) instead of libtorch. The WGPU backend is unstable for SD but may work well in the future as burn-wpu is optimized. ```bash -# torch (at least 6 GB VRAM, possibly less) -# Arguments: [cuda, mps, cpu] +# WGPU/Vulkan backend (GPU accelerated, requires Vulkan-compatible GPU) +# Arguments: -# Cuda -cargo run --release --bin sample burn SDv1-4 7.5 20 "An ancient mossy stone." img cuda - -# Mps(Mac) -cargo run --release --bin sample burn SDv1-4 7.5 20 "An ancient mossy stone." img mps - -# wgpu (UNSTABLE) -# Arguments: +# GPU (Vulkan) cargo run --release --features wgpu-backend --bin sample burn SDv1-4 7.5 20 "An ancient mossy stone." img -``` + +# CPU (UNSTABLE - fallback if GPU not available) +cargo run --release --bin sample burn SDv1-4 7.5 20 "An ancient mossy stone." img This command will generate an image according to the provided prompt, which will be saved as 'img0.png'. diff --git a/src/bin/sample/main.rs b/src/bin/sample/main.rs index e5a146c..56acd8c 100644 --- a/src/bin/sample/main.rs +++ b/src/bin/sample/main.rs @@ -14,7 +14,7 @@ cfg_if::cfg_if! { if #[cfg(feature = "wgpu-backend")] { use burn_wgpu::{Wgpu, WgpuDevice}; } else { - use burn_tch::{LibTorch, LibTorchDevice}; + use burn_ndarray::NdArrayDevice; } } @@ -61,24 +61,8 @@ fn main() { type Backend = Wgpu; let device = WgpuDevice::BestAvailable; } else { - type Backend = LibTorch; - - let device = if let Some(dev_str) = device_arg { - match dev_str.to_lowercase().as_str() { - "cpu" => LibTorchDevice::Cpu, - "mps" => LibTorchDevice::Mps, - s if s.starts_with("cuda") => { - let idx = s[4..].parse().unwrap_or(0); - LibTorchDevice::Cuda(idx) - } - _ => { - eprintln!("Unknown device: {}", dev_str); - process::exit(1); - } - } - } else { - LibTorchDevice::Cuda(0) - }; + type Backend = burn::backend::ndarray::NdArray; + let device = NdArrayDevice::Cpu; } }