huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Enable stable-diffusion 3 on metal. (#2560)

Browse Source
This commit is contained in:
Laurent Mazare
2024-10-14 08:59:12 +02:00
committed by GitHub
parent f553ab5eb4
commit 3d1dc06cdb
4 changed files with 11 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
candle-examples/examples/stable-diffusion-3/main.rs
View File

@ -30,9 +30,9 @@ struct Args {
#[arg(long)]
cpu: bool,
/// The CUDA device ID to use.
#[arg(long, default_value = "0")]
cuda_device_id: usize,
/// The GPU device ID to use.
#[arg(long, default_value_t = 0)]
gpu_device_id: usize,
/// Enable tracing (generates a trace-timestamp.json file).
#[arg(long)]
@ -81,7 +81,7 @@ fn run(args: Args) -> Result<()> {
prompt,
uncond_prompt,
cpu,
cuda_device_id,
gpu_device_id,
tracing,
use_flash_attn,
height,
@ -100,11 +100,14 @@ fn run(args: Args) -> Result<()> {
None
};
// TODO: Support and test on Metal.
let device = if cpu {
candle::Device::Cpu
} else if candle::utils::cuda_is_available() {
candle::Device::new_cuda(gpu_device_id)?
} else if candle::utils::metal_is_available() {
candle::Device::new_metal(gpu_device_id)?
} else {
candle::Device::cuda_if_available(cuda_device_id)?
candle::Device::Cpu
};
let api = hf_hub::api::sync::Api::new()?;