Write the generated images using the image crate. (huggingface#363)

* Use the image crate to write the generated images. * Make the dependency optional.
philparzer · Aug 9, 2023 · 3a62aee · 3a62aee
1 parent be21d7e
commit 3a62aee
Showing 4 changed files with 31 additions and 6 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -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
@@ -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
@@ -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
@@ -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(())
+}