consistency decoder

[williamberman]

from diffusers import StableDiffusionPipeline, ConsistencyDecoderVae
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
)
pipe.to('cuda')
pipe("horse", generator=torch.Generator('cpu').manual_seed(0)).images[0].save('orig.png')

vae = ConsistencyDecoderVae.from_pretrained("../consistency-decoder")
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, vae=vae
)
pipe.to('cuda')

pipe("horse", generator=torch.Generator('cpu').manual_seed(0)).images[0].save('new.png')

orig

new

import torch
from diffusers import ConsistencyDecoderVae, AutoencoderKL
from diffusers.utils.loading_utils import load_image
import numpy as np
from PIL import Image

vae_orig = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, subfolder="vae")
vae_orig.to('cuda')

vae_new = ConsistencyDecoderVae.from_pretrained("../consistency-decoder", torch_dtype=torch.float16)
vae_new.to('cuda')


image = load_image("https://raw.githubusercontent.com/openai/consistencydecoder/main/assets/gt1.png").resize((256, 256))
image = torch.from_numpy(np.array(image).transpose(2, 0, 1).astype(np.float32) / 127.5 - 1)[None, :, :, :].half().cuda()
latent = vae_new.encode(image).latent_dist.mean

sample_orig = vae_orig.decode(latent).sample
sample_new = vae_new.decode(latent).sample

Image.fromarray(((sample_orig.detach()[0].cpu().numpy().transpose(1, 2, 0) + 1) * 127.5).clip(0, 255).astype(np.uint8)).save('orig.png')
Image.fromarray(((sample_new.detach()[0].cpu().numpy().transpose(1, 2, 0) + 1) * 127.5).clip(0, 255).astype(np.uint8)).save('new.png')

original

new

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[williamberman]

Need to implement more timestep options but I'm ok shipping w/out and adding a better error message for now

[sayakpaul]

@williamberman summary of the changes:

• Docs
• Tests (including a slow one)

[DN6]

LGTM 👍🏽 Had some comments on a few small things.

[patrickvonplaten]

Very nice job on porting the checkpoint this quickly @williamberman .

I feel quite strongly that we should allow this decoder to be a drop-in replacement for existing SD15 pipelines. E.g. the following should work:

+from diffusers import StableDiffusionPipeline, ConsistencyDecoderVae
-from diffusers import StableDiffusionPipeline
import torch

+vae = ConsistencyDecoderVae.from_pretrained("openai/consistency_decoder", torch_dtype=torch.float16)

model_id = "runwayml/stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, vae=vae)
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]  
    
image.save("astronaut_rides_horse.png")

Therefore, I would do the following:

• 1. Instead create a ConsistencyDecoderVae class that has exactly the same signatures as AutoencoderKL
• 2. Save both the unet, scheduler and encoder (same as the SDv15 vae) as class attributes
• 3. Add the sdv15 vae to the checkpoint as well so that Img2Img pipelines work

class ConsistencyDecoderVae(ModelMixin, ConfigMixin):
    
    @register_to_config
    def __init__(self, <all_config_params_of_unet_and_vae_encoder>):
         # same encoder as standard VAE
         self.encoder = Encoder(
            in_channels=in_channels,
            out_channels=latent_channels,
            down_block_types=down_block_types,
            block_out_channels=block_out_channels,
            layers_per_block=layers_per_block,
            act_fn=act_fn,
            norm_num_groups=norm_num_groups,
            double_z=True,
         )

        self.decoder_unet = UNet2DModel(<all_config_params_on_unet>)
        self.decoder_scheduler = ConsistencyDecoderScheduler()

    def forward(
        self,
        sample: torch.FloatTensor,
        sample_posterior: bool = False,
        return_dict: bool = True,
        generator: Optional[torch.Generator] = None,
    ) -> Union[DecoderOutput, torch.FloatTensor]:
        r"""
        Args:
            sample (`torch.FloatTensor`): Input sample.
            sample_posterior (`bool`, *optional*, defaults to `False`):
                Whether to sample from the posterior.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
        """
        x = sample
        posterior = self.encode(x).latent_dist
        if sample_posterior:
            z = posterior.sample(generator=generator)
        else:
            z = posterior.mode()
        dec = self.decode(z).sample

        if not return_dict:
            return (dec,)

        return DecoderOutput(sample=dec)

    @apply_forward_hook
    def encode(
        self, x: torch.FloatTensor, return_dict: bool = True
    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
        """
        Encode a batch of images into latents.

        Args:
            x (`torch.FloatTensor`): Input batch of images.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.

        Returns:
                The latent representations of the encoded images. If `return_dict` is True, a
                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
        """
        if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
            return self.tiled_encode(x, return_dict=return_dict)

        if self.use_slicing and x.shape[0] > 1:
            encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)]
            h = torch.cat(encoded_slices)
        else:
            h = self.encoder(x)

        moments = self.quant_conv(h)
        posterior = DiagonalGaussianDistribution(moments)

        if not return_dict:
            return (posterior,)

        return AutoencoderKLOutput(latent_dist=posterior)

    @apply_forward_hook
    def decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]:
        # here call the scheduler and the unet ...

For now let's maybe just not allow saving the and loading the scheduler and always default to the correct default values until we have more usage (we can always add it later IMO)

[sayakpaul]

But @patrickvonplaten won't that be restrictive because users won't be able to control the number of timesteps for the decoder scheduler (I know it's currently not supported).

Do you think most users would stick to 2 timesteps? I think they would.

[williamberman]

yeah chatted with patrick here, I have no strong opinions on pipeline vs model class

[williamberman]

This helps for tests for the noise generation in the decoder, I also added an un-used generator arg to the AutoencoderKL.decode so should be fine

[patrickvonplaten]

Perf!

[williamberman]

Helps keep interface same with consistency decoder and can pass generator through from pipeline

[patrickvonplaten]

Yes perf

[williamberman]

Needed to pass generator with same seed for multiple forward calls

[williamberman]

failing tests are unrelated

[patrickvonplaten]

Some bug fixes here: #5722