From 80f0482f26cfa2e595e413576090690c9e23a762 Mon Sep 17 00:00:00 2001 From: Laurent Mazare Date: Thu, 10 Aug 2023 19:24:31 +0200 Subject: [PATCH] Fix the stable-diffusion vae. (#398) * Fix the stable-diffusion vae. * Fix for saving images. --- candle-examples/examples/stable-diffusion/attention.rs | 2 +- candle-examples/examples/stable-diffusion/main.rs | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/candle-examples/examples/stable-diffusion/attention.rs b/candle-examples/examples/stable-diffusion/attention.rs index 83e7ef34..943a6ae1 100644 --- a/candle-examples/examples/stable-diffusion/attention.rs +++ b/candle-examples/examples/stable-diffusion/attention.rs @@ -432,7 +432,7 @@ impl AttentionBlock { (query_states * scale)?.matmul(&(key_states.t()? * scale)?)?; let attention_probs = nn::ops::softmax(&attention_scores, D::Minus1)?; - let xs = attention_probs.matmul(&value_states)?; + let xs = attention_probs.matmul(&value_states.contiguous()?)?; let xs = xs.transpose(1, 2)?.contiguous()?; let xs = xs.flatten_from(D::Minus2)?; let xs = self 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)? }