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

Fix the stable-diffusion vae. (#398)

Browse Source 
* Fix the stable-diffusion vae.

* Fix for saving images.
This commit is contained in:
Laurent Mazare
2023-08-10 19:24:31 +02:00
committed by GitHub
parent 94eff56aee
commit 80f0482f26
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
candle-examples/examples/stable-diffusion/attention.rs
View File

@ -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

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

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