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

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

Browse Source 
* Use the image crate to write the generated images.

* Make the dependency optional.
This commit is contained in:
Laurent Mazare
2023-08-09 16:26:44 +02:00
committed by GitHub
parent be21d7e75a
commit 3a62aee91f
4 changed files with 31 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
candle-examples/examples/stable-diffusion/main.rs
View File

@ -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(())
}

19
candle-examples/examples/stable-diffusion/utils.rs
View File

@ -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(())
}