feat: migrate to pure Rust backend with WGPU/Vulkan support

- Remove libtorch dependency and related features
- Switch to ndarray backend as default for CPU execution
- Update README to reflect new WGPU/Vulkan backend usage
- Simplify device selection to use only CPU backend
- Enable WGPU backend via feature flag for GPU acceleration

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"

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)
+# WGPU/Vulkan backend (GPU accelerated, requires Vulkan-compatible GPU)
-# Arguments: <model_type(burn or dump)> <model_name> <unconditional_guidance_scale> <n_diffusion_steps> <prompt> <output_image_name> [cuda, mps, cpu]
+# Arguments: <model_type(burn or dump)> <model_name> <unconditional_guidance_scale> <n_diffusion_steps> <prompt> <output_image_name>
 
-# Cuda
+# GPU (Vulkan)
-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: <model_type(burn or dump)> <model> <unconditional_guidance_scale> <n_diffusion_steps> <prompt> <output_image>
 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'.
 

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<f32>;
+            type Backend = burn::backend::ndarray::NdArray<f32>;
-
+            let device = NdArrayDevice::Cpu;
-            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)
-            };
         }
     }