FreeU

README.md

@@ -52,7 +52,9 @@ Gradio app for Stable Diffusion 1.5 including:
 * multiple samplers with Karras schedule
 * Compel prompting
 * 100+ styles from sdxl_prompt_styler
-* DeepCache and ToMe
+* FreeU and Clip Skip for quality
+* DeepCache and ToMe for speed
+* Real-ESRGAN upscaling
 * optional TAESD
 
 ## Usage

app.py

@@ -171,9 +171,19 @@ with gr.Blocks(
                             label="Karras σ",
                             value=True,
                         )
+                        use_freeu = gr.Checkbox(
+                            elem_classes=["checkbox"],
+                            label="FreeU",
+                            value=False,
+                        )
+                        use_clip_skip = gr.Checkbox(
+                            elem_classes=["checkbox"],
+                            label="Clip skip",
+                            value=False,
+                        )
                         increment_seed = gr.Checkbox(
                             elem_classes=["checkbox"],
-                            label="Autoincrement",
+                            label="Autoincrement seed",
                             value=True,
                         )
 
@@ -207,11 +217,6 @@ with gr.Blocks(
                             label="Tiny VAE",
                             value=False,
                         )
-                        use_clip_skip = gr.Checkbox(
-                            elem_classes=["checkbox"],
-                            label="Clip skip",
-                            value=False,
-                        )
                         truncate_prompts = gr.Checkbox(
                             elem_classes=["checkbox"],
                             label="Truncate prompts",
@@ -296,6 +301,7 @@ with gr.Blocks(
             num_images,
             use_karras,
             use_taesd,
+            use_freeu,
             use_clip_skip,
             truncate_prompts,
             increment_seed,

cli.py

@@ -34,6 +34,7 @@ def main():
     parser.add_argument("--clip-skip", action="store_true")
     parser.add_argument("--truncate", action="store_true")
     parser.add_argument("--karras", action="store_true")
+    parser.add_argument("--freeu", action="store_true")
     parser.add_argument("--no-increment", action="store_false")
     # fmt: on
 
@@ -52,6 +53,7 @@ def main():
         args.images,
         args.karras,
         args.taesd,
+        args.freeu,
         args.clip_skip,
         args.truncate,
         args.no_increment,

lib/inference.py

@@ -86,6 +86,7 @@ def generate(
     num_images=1,
     karras=False,
     taesd=False,
+    freeu=False,
     clip_skip=False,
     truncate_prompts=False,
     increment_seed=True,
@@ -124,6 +125,7 @@ def generate(
             scheduler,
             karras,
             taesd,
+            freeu,
             deepcache_interval,
             scale,
             DTYPE,

lib/loader.py

@@ -64,6 +64,17 @@ class Loader:
         self.pipe.deepcache.set_params(cache_interval=interval)
         self.pipe.deepcache.enable()
 
+    def _load_freeu(self, freeu=False):
+        # https://github.com/huggingface/diffusers/blob/v0.30.0/src/diffusers/models/unets/unet_2d_condition.py
+        block = self.pipe.unet.up_blocks[0]
+        attrs = ["b1", "b2", "s1", "s2"]
+        has_freeu = all(getattr(block, attr, None) is not None for attr in attrs)
+        if has_freeu and not freeu:
+            self.pipe.disable_freeu()
+        elif not has_freeu and freeu:
+            # https://github.com/ChenyangSi/FreeU
+            self.pipe.enable_freeu(b1=1.5, b2=1.6, s1=0.9, s2=0.2)
+
     def _load_vae(self, model_name=None, taesd=False, variant=None):
         vae_type = type(self.pipe.vae)
         is_kl = issubclass(vae_type, (AutoencoderKL, OptimizedModule))
@@ -93,7 +104,18 @@ class Loader:
                 model=model,
             )
 
-    def load(self, model, scheduler, karras, taesd, deepcache_interval, scale, dtype, device):
+    def load(
+        self,
+        model,
+        scheduler,
+        karras,
+        taesd,
+        freeu,
+        deepcache_interval,
+        scale,
+        dtype,
+        device,
+    ):
         model_lower = model.lower()
 
         schedulers = {
@@ -155,8 +177,9 @@ class Loader:
                 if not same_scheduler or not same_karras:
                     self.pipe.scheduler = schedulers[scheduler](**scheduler_kwargs)
                 self._load_vae(model_lower, taesd, variant)
-                self._load_deepcache(interval=deepcache_interval)
-                self._load_upscaler(device=device, scale=scale)
+                self._load_freeu(freeu)
+                self._load_deepcache(deepcache_interval)
+                self._load_upscaler(device, scale)
                 torch.cuda.empty_cache()
                 return self.pipe, self.upscaler
             else:
@@ -173,7 +196,8 @@ class Loader:
             tokens=list(EMBEDDINGS.values()),
         )
         self._load_vae(model_lower, taesd, variant)
-        self._load_deepcache(interval=deepcache_interval)
-        self._load_upscaler(device=device, scale=scale)
+        self._load_freeu(freeu)
+        self._load_deepcache(deepcache_interval)
+        self._load_upscaler(device, scale)
         torch.cuda.empty_cache()
         return self.pipe, self.upscaler

usage.md

@@ -32,6 +32,14 @@ Arrays allow you to generate different images from a single prompt. For example,
 
 Styles are prompt templates from twri's [sdxl_prompt_styler](https://github.com/twri/sdxl_prompt_styler) Comfy node. Start with a subject like "cat", pick a style, and iterate from there.
 
+#### FreeU
+
+[FreeU](https://github.com/ChenyangSi/FreeU) (Si et al. 2023) re-weights the contributions sourced from the U-Net’s skip connections and backbone feature maps to potentially improve image quality.
+
+#### Clip Skip
+
+When enabled, the last CLIP layer is skipped. This _can_ improve image quality with anime models.
+
 ### Scale
 
 Rescale up to 4x using [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN).
@@ -63,7 +71,7 @@ Optionally, the [Karras](https://arxiv.org/abs/2206.00364) noise schedule can be
 
 #### DeepCache
 
-[DeepCache](https://github.com/horseee/DeepCache) (Ma et al. 2023) caches lower UNet layers and reuses them every `Interval` steps:
+[DeepCache](https://github.com/horseee/DeepCache) (Ma et al. 2023) caches lower U-Net layers and reuses them every `Interval` steps:
 * `1`: no caching
 * `2`: more quality (default)
 * `3`: balanced
@@ -71,16 +79,12 @@ Optionally, the [Karras](https://arxiv.org/abs/2206.00364) noise schedule can be
 
 #### ToMe
 
-[Token merging](https://arxiv.org/abs/2303.17604) (Bolya & Hoffman 2023) reduces the number of tokens processed by the model. Set `Ratio` to the desired reduction factor. ToMe's impact is more noticeable on larger images.
+[Token merging](https://github.com/dbolya/tomesd) (Bolya & Hoffman 2023) reduces the number of tokens processed by the model. Set `Ratio` to the desired reduction factor. ToMe's impact is more noticeable on larger images.
 
 #### Tiny VAE
 
 Enable [madebyollin/taesd](https://github.com/madebyollin/taesd) for almost instant latent decoding with a minor loss in detail. Useful for development.
 
-#### Clip Skip
-
-When enabled, the last CLIP layer is skipped. This _can_ improve image quality with anime models.
-
 #### Prompt Truncation
 
 When enabled, prompts will be truncated to CLIP's limit of 77 tokens. By default this is _disabled_, so Compel will chunk prompts into segments rather than cutting them off.