diff options
Diffstat (limited to 'candle-examples/examples/stable-diffusion/main.rs')
| -rw-r--r-- | candle-examples/examples/stable-diffusion/main.rs | 6 |
1 files changed, 3 insertions, 3 deletions
diff --git a/candle-examples/examples/stable-diffusion/main.rs b/candle-examples/examples/stable-diffusion/main.rs index 8ce0c234..ac31e855 100644 --- a/candle-examples/examples/stable-diffusion/main.rs +++ b/candle-examples/examples/stable-diffusion/main.rs @@ -14,7 +14,7 @@ mod utils; mod vae; use anyhow::{Error as E, Result}; -use candle::{DType, Device, Tensor}; +use candle::{DType, Device, IndexOp, Tensor}; use clap::Parser; use tokenizers::Tokenizer; @@ -245,7 +245,7 @@ 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 = (image * 255.)?.to_dtype(DType::U8)?.i(0)?; let image_filename = output_filename(&final_image, idx + 1, num_samples, Some(timestep_index + 1)); crate::utils::save_image(&image, image_filename)? @@ -260,7 +260,7 @@ 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 = (image * 255.)?.to_dtype(DType::U8)?.i(0)?; let image_filename = output_filename(&final_image, idx + 1, num_samples, None); crate::utils::save_image(&image, image_filename)? } |