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

candle-examples/Cargo.toml

@@ -122,6 +122,3 @@ required-features = ["onnx"]
 [[example]]
 name = "colpali"
 required-features = ["pdf2image"]
-
-[[example]]
-name = "stable-diffusion-3"

candle-examples/examples/stable-diffusion-3/main.rs

@@ -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()?;

candle-examples/examples/stable-diffusion-3/sampling.rs

@@ -31,7 +31,7 @@ pub fn euler_sample(
  let timestep = (*s_curr) * 1000.0;
  let noise_pred = mmdit.forward(
  &Tensor::cat(&[x.clone(), x.clone()], 0)?,
- &Tensor::full(timestep, (2,), x.device())?.contiguous()?,
+ &Tensor::full(timestep as f32, (2,), x.device())?.contiguous()?,
  y,
  context,
  )?;

candle-transformers/src/models/marian.rs

@@ -1,9 +1,8 @@
 use super::with_tracing::{linear, Embedding, Linear};
 use candle::{Result, Tensor};
 use candle_nn::{layer_norm, LayerNorm, VarBuilder};
-use serde::Deserialize;
 
-#[derive(Debug, Clone, Deserialize)]
+#[derive(Debug, Clone, serde::Deserialize)]
 pub struct Config {
  pub vocab_size: usize,
  pub decoder_vocab_size: Option<usize>,