diff options
| -rw-r--r-- | Cargo.toml | 1 | ||||
| -rw-r--r-- | candle-examples/Cargo.toml | 5 | ||||
| -rw-r--r-- | candle-examples/examples/stable-diffusion/main.rs | 12 | ||||
| -rw-r--r-- | candle-examples/examples/stable-diffusion/utils.rs | 19 |
4 files changed, 31 insertions, 6 deletions
@@ -33,6 +33,7 @@ cudarc = { version = "0.9.13", features = ["f16"] } gemm = { version = "0.15.5", package = "candle-gemm" } hf-hub = "0.2.0" half = { version = "2.3.1", features = ["num-traits", "rand_distr"] } +image = { version = "0.24.7", default-features = false, features = ["jpeg", "png"] } intel-mkl-src = { version = "0.8.1", features = ["mkl-static-lp64-iomp"] } libc = { version = "0.2.147" } log = "0.4" diff --git a/candle-examples/Cargo.toml b/candle-examples/Cargo.toml index f3a4e325..84115dd5 100644 --- a/candle-examples/Cargo.toml +++ b/candle-examples/Cargo.toml @@ -23,6 +23,7 @@ num-traits = { workspace = true } intel-mkl-src = { workspace = true, optional = true } cudarc = { workspace = true, optional = true } half = { workspace = true, optional = true } +image = { workspace = true, optional = true } [dev-dependencies] anyhow = { workspace = true } @@ -51,3 +52,7 @@ nccl = ["cuda", "cudarc/nccl", "dep:half"] [[example]] name = "llama_multiprocess" required-features = ["cuda", "nccl", "flash-attn"] + +[[example]] +name = "stable-diffusion" +required-features = ["image"] diff --git a/candle-examples/examples/stable-diffusion/main.rs b/candle-examples/examples/stable-diffusion/main.rs index 8bb3c56d..8ce0c234 100644 --- a/candle-examples/examples/stable-diffusion/main.rs +++ b/candle-examples/examples/stable-diffusion/main.rs @@ -245,10 +245,10 @@ fn run(args: Args) -> Result<()> { if args.intermediary_images { let image = vae.decode(&(&latents / 0.18215)?)?; let image = ((image / 2.)? + 0.5)?.to_device(&Device::Cpu)?; - let _image = (image * 255.)?.to_dtype(DType::U8); - let _image_filename = + let image = (image * 255.)?.to_dtype(DType::U8)?; + let image_filename = output_filename(&final_image, idx + 1, num_samples, Some(timestep_index + 1)); - // TODO: save igame + crate::utils::save_image(&image, image_filename)? } } @@ -260,9 +260,9 @@ fn run(args: Args) -> Result<()> { let image = vae.decode(&(&latents / 0.18215)?)?; // TODO: Add the clamping between 0 and 1. let image = ((image / 2.)? + 0.5)?.to_device(&Device::Cpu)?; - let _image = (image * 255.)?.to_dtype(DType::U8); - let _image_filename = output_filename(&final_image, idx + 1, num_samples, None); - // TODO: save image. + let image = (image * 255.)?.to_dtype(DType::U8)?; + let image_filename = output_filename(&final_image, idx + 1, num_samples, None); + crate::utils::save_image(&image, image_filename)? } Ok(()) } diff --git a/candle-examples/examples/stable-diffusion/utils.rs b/candle-examples/examples/stable-diffusion/utils.rs index 0c95cfef..ef4dd956 100644 --- a/candle-examples/examples/stable-diffusion/utils.rs +++ b/candle-examples/examples/stable-diffusion/utils.rs @@ -10,3 +10,22 @@ pub fn linspace(start: f64, stop: f64, steps: usize) -> Result<Tensor> { .collect::<Vec<_>>(); Tensor::from_vec(vs, steps, &Device::Cpu) } + +/// Saves an image to disk using the image crate, this expects an input with shape +/// (c, width, height). +pub fn save_image<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<()> { + let p = p.as_ref(); + let (channel, width, height) = img.dims3()?; + if channel != 3 { + candle::bail!("save_image expects an input of shape (3, width, height)") + } + let img = img.transpose(0, 1)?.t()?.flatten_all()?; + let pixels = img.to_vec1::<u8>()?; + let image: image::ImageBuffer<image::Rgb<u8>, Vec<u8>> = + match image::ImageBuffer::from_raw(width as u32, height as u32, pixels) { + Some(image) => image, + None => candle::bail!("error saving image {p:?}"), + }; + image.save(p).map_err(candle::Error::wrap)?; + Ok(()) +} |