Support OpenAI's Consistency Decoder

[Pirog17000]

Source - https://github.com/openai/consistencydecoder

OpenAI open-sourced consistency decoder, that decodes latent like a pro. Makes images more consistent and less disturbing

But. Right now I see no options to properly decode latent with their provided code. So, maybe there's some kind of solution already exist, but seems that it might be something you can add in diffusers as official support.
Weights of the Consistency Decoder are in .pt format. Possibly it can be converted into .safetensors and loaded as regular VAE.
Also, it is memory-heavy, meaning to decode image quick there's real need to unload the pipeline/model from VRAM at least to cpu

Here's my code which results into garbage-grade images (specifically out of SDXL pipeline):

import consistencydecoder
from diffusers import StableDiffusionXLPipeline
import torch
from consistencydecoder import ConsistencyDecoder

seed = 42
generator = torch.Generator(device="cuda").manual_seed(seed)

print(f"Loading pipe")
pipe = StableDiffusionXLPipeline.from_single_file(
    model_path, torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

pipe.vae.cuda()
decoder_consistency = ConsistencyDecoder(device="cuda:0")

image = pipe(prompt=prompt, width=1024, height=1024,
             num_inference_steps=20,
             generator=generator,
             output_type="latent",
             ).images

pipe.to("cpu")
print(f"Decoding latent")
consistencydecoder.save_image(image, "output.png")

[fhlt]

rescale with 0.18215 is all you need
image = 1.0 / 0.18215 * image
image = decoder_consistency(image)

[zzzzzero]

consistencydecoder is for SD 1.5(2.0),not for SDXL,since SDXL uses a different VAE;
https://www.reddit.com/r/StableDiffusion/comments/17pal90/comment/k850nyk/?utm_source=share&utm_medium=web3x&utm_name=web3xcss&utm_term=1&utm_content=share_button

Here is my colab link:
https://colab.research.google.com/drive/1uCJLGU17c-ytcelWUXdHOKpAq_S2SR_e?usp=sharing
TL,DR ，The consistency decoder can achieve better results than the VAE decoder in subtle areas of some images， and the consistency decoder takes much longer time.
Just use consistency decoder like this:

import torch
from diffusers import StableDiffusionPipeline
from consistencydecoder import ConsistencyDecoder, save_image, load_image

# encode with stable diffusion vae
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, device="cuda:0"
)
pipe = pipe.to("cuda")
decoder_consistency = ConsistencyDecoder(device="cuda:0") # Model size: 2.49 GB

output = pipe(prompt="wonderful building", width=512, height=512, output_type="latent")
latents = output.images[0]
prepped = latents.unsqueeze(0) / pipe.vae.config.scaling_factor

#sd1.5 vae decoder
vae_img =pipe.vae.decode(prepped).sample.detach()
vae_image=pipe.image_processor.postprocess(vae_img,output_type='pil',do_denormalize=[True])
vae_image[0]


#consisency decoder
img = decoder_consistency(prepped)
image=pipe.image_processor.postprocess(img,output_type='pil',do_denormalize=[True])
image[0]

[dg845]

Would be interested in working on this if the maintainers think it's a good idea :).

[patrickvonplaten]

#5694