use tinyVAE

app-img2img.py

@@ -55,10 +55,9 @@ else:
         custom_pipeline="latent_consistency_img2img.py",
         custom_revision="main",
     )
-# TODO try to use tiny VAE
-# pipe.vae = AutoencoderTiny.from_pretrained(
-#     "madebyollin/taesd", torch_dtype=torch.float16, use_safetensors=True
-# )
+pipe.vae = AutoencoderTiny.from_pretrained(
+    "madebyollin/taesd", torch_dtype=torch.float16, use_safetensors=True
+)
 pipe.set_progress_bar_config(disable=True)
 pipe.to(torch_device=torch_device, torch_dtype=torch_dtype).to(device)
 pipe.unet.to(memory_format=torch.channels_last)

latent_consistency_img2img.py

@@ -25,6 +25,7 @@ import torch
 from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
 
 from diffusers import (
+    AutoencoderTiny,
     AutoencoderKL,
     ConfigMixin,
     DiffusionPipeline,
@@ -226,13 +227,22 @@ class LatentConsistencyModelImg2ImgPipeline(DiffusionPipeline):
                 )
 
             elif isinstance(generator, list):
-                init_latents = [
-                    self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i])
-                    for i in range(batch_size)
-                ]
+                if isinstance(self.vae, AutoencoderTiny):
+                    init_latents = [
+                        self.vae.encode(image[i : i + 1]).latents
+                        for i in range(batch_size)
+                    ]
+                else:
+                    init_latents = [
+                        self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i])
+                        for i in range(batch_size)
+                    ]
                 init_latents = torch.cat(init_latents, dim=0)
             else:
-                init_latents = self.vae.encode(image).latent_dist.sample(generator)
+                if isinstance(self.vae, AutoencoderTiny):
+                    init_latents = self.vae.encode(image).latents
+                else:
+                    init_latents = self.vae.encode(image).latent_dist.sample(generator)
 
             init_latents = self.vae.config.scaling_factor * init_latents