Remove tiny VAE

DOCS.md

@@ -95,7 +95,3 @@ Enable `Use negative TI` to append [`fast_negative`](https://civitai.com/models/
 #### Clip Skip
 
 When enabled, the last CLIP layer is skipped. Can sometimes improve image quality.
-
-#### Tiny VAE
-
-Enable [madebyollin/taesd](https://github.com/madebyollin/taesd) for near-instant latent decoding with a minor loss in detail. Useful for development.

README.md

@@ -60,14 +60,13 @@ preload_from_hub:  # up to 10
 # diffusion
 
 Gradio app for Stable Diffusion 1.5 featuring:
-* txt2img and img2img pipelines with ControlNet and IP-Adapter
-* Canny edge detection (more preprocessors coming soon)
+* txt2img and img2img pipelines with IP-Adapter
+* ControlNet with Canny edge detection (more preprocessors coming soon)
 * Compel prompt weighting
 * Hand-written style templates
 * Multiple samplers with Karras scheduling
 * DeepCache, FreeU, and Clip Skip available
 * Real-ESRGAN upscaling
-* Optional tiny autoencoder
 
 ## Usage
 

app.py

@@ -302,11 +302,6 @@ with gr.Blocks(
                     label="Use negative TI",
                     value=False,
                 )
-                use_taesd = gr.Checkbox(
-                    elem_classes=["checkbox"],
-                    label="Tiny VAE",
-                    value=False,
-                )
                 use_freeu = gr.Checkbox(
                     elem_classes=["checkbox"],
                     label="FreeU",
@@ -456,7 +451,6 @@ with gr.Blocks(
             scale,
             num_images,
             use_karras,
-            use_taesd,
             use_freeu,
             use_clip_skip,
             use_ip_face,

lib/config.py

@@ -68,7 +68,6 @@ Config = SimpleNamespace(
         "Linaqruf/anything-v3-1": ["anything-v3-2.safetensors"],
         "lllyasviel/control_v11p_sd15_canny": ["diffusion_pytorch_model.fp16.safetensors"],
         "Lykon/dreamshaper-8": [*sd_files],
-        "madebyollin/taesd": ["diffusion_pytorch_model.safetensors"],
         "prompthero/openjourney-v4": ["openjourney-v4.ckpt"],
         "SG161222/Realistic_Vision_V5.1_noVAE": ["Realistic_Vision_V5.1_fp16-no-ema.safetensors"],
         "XpucT/Deliberate": ["Deliberate_v6.safetensors"],

lib/inference.py

@@ -80,7 +80,6 @@ def generate(
     scale=1,
     num_images=1,
     karras=False,
-    taesd=False,
     freeu=False,
     clip_skip=False,
     ip_face=False,
@@ -144,7 +143,6 @@ def generate(
         deepcache,
         scale,
         karras,
-        taesd,
         freeu,
         progress,
     )

lib/loader.py

@@ -4,7 +4,7 @@ from threading import Lock
 import torch
 from DeepCache import DeepCacheSDHelper
 from diffusers import ControlNetModel
-from diffusers.models import AutoencoderKL, AutoencoderTiny
+from diffusers.models import AutoencoderKL
 from diffusers.models.attention_processor import AttnProcessor2_0, IPAdapterAttnProcessor2_0
 
 from .config import Config
@@ -29,20 +29,6 @@ class Loader:
                 cls._instance.log = Logger("Loader")
         return cls._instance
 
-    @property
-    def _is_kl_vae(self):
-        if self.pipe is not None:
-            vae_type = type(self.pipe.vae)
-            return issubclass(vae_type, AutoencoderKL)
-        return False
-
-    @property
-    def _is_tiny_vae(self):
-        if self.pipe is not None:
-            vae_type = type(self.pipe.vae)
-            return issubclass(vae_type, AutoencoderTiny)
-        return False
-
     @property
     def _has_freeu(self):
         if self.pipe is not None:
@@ -184,6 +170,7 @@ class Loader:
             to_unload.append("model")
             to_unload.append("pipe")
 
+        # Flush cache and run garbage collector
         clear_cuda_cache()
         for component in to_unload:
             setattr(self, component, None)
@@ -284,32 +271,22 @@ class Loader:
         if self.pipe is not None:
             self.pipe.set_progress_bar_config(disable=progress is not None)
 
-    def _load_vae(self, taesd=False, model=""):
-        # by default all models use KL
-        if self._is_kl_vae and taesd:
-            msg = "Loading Tiny VAE"
-            with timer(msg, logger=self.log.info):
-                self.pipe.vae = AutoencoderTiny.from_pretrained(
-                    pretrained_model_name_or_path="madebyollin/taesd",
+    # Handle single-file and diffusers-style models
+    def _load_vae(self, model=""):
+        msg = "Loading VAE"
+        with timer(msg, logger=self.log.info):
+            if model.lower() in Config.MODEL_CHECKPOINTS.keys():
+                self.pipe.vae = AutoencoderKL.from_single_file(
+                    f"https://huggingface.co/{model}/{Config.MODEL_CHECKPOINTS[model.lower()]}",
                     torch_dtype=self.pipe.dtype,
                 ).to(self.pipe.device)
-            return
-
-        if self._is_tiny_vae and not taesd:
-            msg = "Loading KL VAE"
-            with timer(msg, logger=self.log.info):
-                if model.lower() in Config.MODEL_CHECKPOINTS.keys():
-                    self.pipe.vae = AutoencoderKL.from_single_file(
-                        f"https://huggingface.co/{model}/{Config.MODEL_CHECKPOINTS[model.lower()]}",
-                        torch_dtype=self.pipe.dtype,
-                    ).to(self.pipe.device)
-                else:
-                    self.pipe.vae = AutoencoderKL.from_pretrained(
-                        pretrained_model_name_or_path=model,
-                        torch_dtype=self.pipe.dtype,
-                        subfolder="vae",
-                        variant="fp16",
-                    ).to(self.pipe.device)
+            else:
+                self.pipe.vae = AutoencoderKL.from_pretrained(
+                    pretrained_model_name_or_path=model,
+                    torch_dtype=self.pipe.dtype,
+                    subfolder="vae",
+                    variant="fp16",
+                ).to(self.pipe.device)
 
     def load(
         self,
@@ -321,7 +298,6 @@ class Loader:
         deepcache,
         scale,
         karras,
-        taesd,
         freeu,
         progress,
     ):
@@ -397,11 +373,12 @@ class Loader:
             if not same_scheduler or not same_karras:
                 self.pipe.scheduler = Config.SCHEDULERS[scheduler](**scheduler_kwargs)
 
+        # Load VAE
+        self._load_vae(model)
+
         CURRENT_STEP = 1
         TOTAL_STEPS = sum(
             [
-                self._is_kl_vae and taesd,
-                self._is_tiny_vae and not taesd,
                 self._should_load_freeu(freeu),
                 self._should_load_deepcache(deepcache),
                 self._should_load_ip_adapter(ip_adapter),
@@ -428,8 +405,3 @@ class Loader:
         if self._should_load_upscaler(scale):
             self._load_upscaler(scale)
             safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)
-            CURRENT_STEP += 1
-
-        if self._is_kl_vae and taesd or self._is_tiny_vae and not taesd:
-            self._load_vae(taesd, model)
-            safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)