Add DeepCache and T-GATE

app.css

@@ -28,7 +28,7 @@
   margin-left: 8px;
 }
 
-#gallery {
+.gallery {
   --block-border-width: 0px;
   background-color: transparent;
 }

app.py

@@ -4,9 +4,11 @@ import gradio as gr
 
 from generate import generate
 
+DEFAULT_NEGATIVE_PROMPT = "<bad_prompt>, ugly, unattractive, deformed, disfigured, mutated, blurry, distorted, noisy, grainy, glitch, worst quality"
+
 # base font stacks
-mono_fonts = ["monospace"]
-sans_fonts = [
+MONO_FONTS = ["monospace"]
+SANS_FONTS = [
     "sans-serif",
     "Apple Color Emoji",
     "Segoe UI Emoji",
@@ -42,96 +44,109 @@ def generate_btn_click(*args, **kwargs):
 
 
 with gr.Blocks(
-    head=read_file("head.html"),
+    head=read_file("./partials/head.html"),
     css="./app.css",
     js="./app.js",
     theme=gr.themes.Default(
         # colors
+        neutral_hue=gr.themes.colors.gray,
         primary_hue=gr.themes.colors.orange,
         secondary_hue=gr.themes.colors.blue,
-        neutral_hue=gr.themes.colors.gray,
         # sizing
         text_size=gr.themes.sizes.text_md,
-        spacing_size=gr.themes.sizes.spacing_md,
         radius_size=gr.themes.sizes.radius_sm,
+        spacing_size=gr.themes.sizes.spacing_md,
         # fonts
-        font=[gr.themes.GoogleFont("Inter"), *sans_fonts],
-        font_mono=[gr.themes.GoogleFont("Ubuntu Mono"), *mono_fonts],
+        font=[gr.themes.GoogleFont("Inter"), *SANS_FONTS],
+        font_mono=[gr.themes.GoogleFont("Ubuntu Mono"), *MONO_FONTS],
     ).set(
-        block_background_fill=gr.themes.colors.gray.c50,
-        block_background_fill_dark=gr.themes.colors.gray.c900,
         block_shadow="0 0 #0000",
         block_shadow_dark="0 0 #0000",
+        block_background_fill=gr.themes.colors.gray.c50,
+        block_background_fill_dark=gr.themes.colors.gray.c900,
     ),
 ) as demo:
-    gr.HTML(read_file("intro.html"))
+    gr.HTML(read_file("./partials/intro.html"))
     output_images = gr.Gallery(
-        label="Output",
-        show_label=False,
-        columns=1,
-        interactive=False,
+        elem_classes=["gallery"],
         show_share_button=False,
-        elem_id="gallery",
+        interactive=False,
+        show_label=False,
+        label="Output",
+        format="png",
+        columns=2,
     )
     prompt = gr.Textbox(
-        label="Prompt",
+        placeholder="corgi, at the beach, cute, 8k",
         show_label=False,
-        lines=2,
-        placeholder="corgi, at the beach, cute",
+        label="Prompt",
         value=None,
+        lines=2,
     )
     generate_btn = gr.Button("Generate", variant="primary", elem_classes=[])
 
     with gr.Accordion(
+        elem_classes=["accordion"],
+        elem_id="menu",
         label="Menu",
         open=False,
-        elem_id="menu",
-        elem_classes=["accordion"],
     ):
         with gr.Tabs():
             with gr.TabItem("⚙️ Settings"):
                 with gr.Group():
                     negative_prompt = gr.Textbox(
                         label="Negative Prompt",
+                        value=DEFAULT_NEGATIVE_PROMPT,
+                        placeholder="",
                         lines=1,
-                        placeholder="ugly",
-                        value="",
                     )
 
                     with gr.Row():
                         num_images = gr.Dropdown(
-                            label="Images",
                             choices=[1, 2, 3, 4],
-                            value=1,
                             filterable=False,
+                            label="Images",
+                            value=1,
+                            scale=1,
                         )
-                        aspect_ratio = gr.Dropdown(
-                            label="Aspect Ratio",
-                            choices=["1:1", "4:3", "3:4", "16:9", "9:16"],
-                            value="1:1",
-                            filterable=False,
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=1024,
+                            value=512,
+                            step=32,
+                            scale=2,
+                        )
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=1024,
+                            value=512,
+                            step=32,
+                            scale=2,
                         )
-                        seed = gr.Number(label="Seed", value=0)
 
                     with gr.Row():
                         guidance_scale = gr.Slider(
                             label="Guidance Scale",
                             minimum=1.0,
                             maximum=15.0,
-                            step=0.1,
                             value=7.5,
+                            step=0.1,
                         )
                         inference_steps = gr.Slider(
                             label="Inference Steps",
                             minimum=1,
                             maximum=50,
-                            step=1,
                             value=30,
+                            step=1,
                         )
 
                     with gr.Row():
                         model = gr.Dropdown(
+                            value="Lykon/dreamshaper-8",
                             label="Model",
+                            scale=2,
                             choices=[
                                 "fluently/Fluently-v4",
                                 "Linaqruf/anything-v3-1",
@@ -140,10 +155,12 @@ with gr.Blocks(
                                 "runwayml/stable-diffusion-v1-5",
                                 "SG161222/Realistic_Vision_V5.1_Novae",
                             ],
-                            value="Lykon/dreamshaper-8",
                         )
                         scheduler = gr.Dropdown(
+                            elem_id="scheduler",
                             label="Scheduler",
+                            value="DEIS 2M",
+                            scale=2,
                             choices=[
                                 "DEIS 2M",
                                 "DPM++ 2M",
@@ -153,22 +170,20 @@ with gr.Blocks(
                                 "LMS",
                                 "PNDM",
                             ],
-                            value="DEIS 2M",
-                            elem_id="scheduler",
                         )
+                        seed = gr.Number(label="Seed", value=42)
 
                     with gr.Row():
                         use_karras = gr.Checkbox(
-                            label="Use Karras σ",
-                            value=True,
                             elem_classes=["checkbox"],
-                            scale=2,
+                            label="Karras σ",
+                            value=True,
+                            scale=1,
                         )
                         increment_seed = gr.Checkbox(
-                            label="Autoincrement seed",
-                            value=True,
                             elem_classes=["checkbox"],
-                            elem_id="increment-seed",
+                            label="Autoincrement",
+                            value=True,
                             scale=2,
                         )
                         random_seed_btn = gr.Button(
@@ -179,21 +194,57 @@ with gr.Blocks(
                         )
 
             with gr.TabItem("🛠️ Advanced"):
-                gr.Markdown("_Coming soon..._", elem_classes=["markdown"])
+                with gr.Group():
+                    with gr.Row():
+                        deep_cache_interval = gr.Slider(
+                            label="DeepCache Interval",
+                            minimum=1,
+                            maximum=4,
+                            value=0,
+                            step=1,
+                        )
+                        deep_cache_branch = gr.Slider(
+                            label="DeepCache Branch",
+                            minimum=0,
+                            maximum=3,
+                            value=0,
+                            step=1,
+                        )
+                        tgate_step = gr.Slider(
+                            label="T-GATE Step",
+                            minimum=0,
+                            maximum=50,
+                            value=0,
+                            step=1,
+                        )
+
+                    with gr.Row():
+                        use_taesd = gr.Checkbox(
+                            elem_classes=["checkbox"],
+                            label="Tiny VAE",
+                            value=False,
+                            scale=1,
+                        )
+                        use_clip_skip = gr.Checkbox(
+                            elem_classes=["checkbox"],
+                            label="Clip skip",
+                            value=False,
+                            scale=1,
+                        )
+                        truncate_prompts = gr.Checkbox(
+                            elem_classes=["checkbox"],
+                            label="Truncate prompts",
+                            value=False,
+                            scale=3,
+                        )
 
             with gr.TabItem("ℹ️ Info"):
                 gr.Markdown(read_file("info.md"), elem_classes=["markdown"])
 
-    # change gallery columns when num_images changes
-    num_images.change(
-        lambda n: gr.Gallery(columns=n),
-        inputs=[num_images],
-        outputs=[output_images],
-    )
-
     # update the random seed using JavaScript
     random_seed_btn.click(None, outputs=[seed], js="() => Math.floor(Math.random() * 2**32)")
 
+    # ensure correct argument order
     generate_btn.click(
         generate_btn_click,
         api_name="generate",
@@ -205,12 +256,19 @@ with gr.Blocks(
             seed,
             model,
             scheduler,
-            aspect_ratio,
+            width,
+            height,
             guidance_scale,
             inference_steps,
-            use_karras,
             num_images,
+            use_karras,
+            use_taesd,
+            use_clip_skip,
+            truncate_prompts,
             increment_seed,
+            deep_cache_interval,
+            deep_cache_branch,
+            tgate_step,
         ],
     )
 

generate.py

@@ -1,12 +1,15 @@
 import re
+from contextlib import contextmanager
 from datetime import datetime
 from itertools import product
 from os import environ
+from types import MethodType
 from warnings import filterwarnings
 
 import spaces
 import torch
-from compel import Compel
+from compel import Compel, DiffusersTextualInversionManager, ReturnedEmbeddingsType
+from DeepCache import DeepCacheSDHelper
 from diffusers import (
     DEISMultistepScheduler,
     DPMSolverMultistepScheduler,
@@ -17,18 +20,23 @@ from diffusers import (
     PNDMScheduler,
     StableDiffusionPipeline,
 )
-from diffusers.models import AutoencoderTiny
+from diffusers.models import AutoencoderKL, AutoencoderTiny
+from tgate.SD import tgate as tgate_sd
+from tgate.SD_DeepCache import tgate as tgate_sd_deepcache
+from torch._dynamo import OptimizedModule
 
 ZERO_GPU = (
     environ.get("SPACES_ZERO_GPU", "").lower() == "true"
     or environ.get("SPACES_ZERO_GPU", "") == "1"
 )
 
-TORCH_DTYPE = (
-    torch.bfloat16
-    if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
-    else torch.float16
-)
+EMBEDDINGS = {
+    "./embeddings/bad_prompt_version2.pt": "<bad_prompt>",
+    "./embeddings/BadDream.pt": "<bad_dream>",
+    "./embeddings/FastNegativeV2.pt": "<fast_negative>",
+    "./embeddings/negative_hand.pt": "<negative_hand>",
+    "./embeddings/UnrealisticDream.pt": "<unrealistic_dream>",
+}
 
 # some models use the deprecated CLIPFeatureExtractor class
 # should use CLIPImageProcessor instead
@@ -46,7 +54,27 @@ class Loader:
             cls._instance.pipe = None
         return cls._instance
 
-    def load(self, model, scheduler, karras):
+    def _load_vae(self, model_name=None, taesd=False, dtype=None):
+        if taesd:
+            # can't compile tiny VAE
+            return AutoencoderTiny.from_pretrained(
+                pretrained_model_name_or_path="madebyollin/taesd",
+                use_safetensors=True,
+                torch_dtype=dtype,
+            ).to(self.gpu)
+
+        return torch.compile(
+            fullgraph=True,
+            mode="reduce-overhead",
+            model=AutoencoderKL.from_pretrained(
+                pretrained_model_name_or_path=model_name,
+                use_safetensors=True,
+                torch_dtype=dtype,
+                subfolder="vae",
+            ).to(self.gpu),
+        )
+
+    def load(self, model, scheduler, karras, taesd, dtype=None):
         model_lower = model.lower()
 
         schedulers = {
@@ -60,24 +88,24 @@ class Loader:
         }
 
         scheduler_kwargs = {
-            "beta_start": 0.00085,
-            "beta_end": 0.012,
             "beta_schedule": "scaled_linear",
             "timestep_spacing": "leading",
-            "steps_offset": 1,
             "use_karras_sigmas": karras,
+            "beta_start": 0.00085,
+            "beta_end": 0.012,
+            "steps_offset": 1,
         }
 
         if scheduler == "PNDM" or scheduler == "Euler a":
             del scheduler_kwargs["use_karras_sigmas"]
 
         pipe_kwargs = {
+            "scheduler": schedulers[scheduler](**scheduler_kwargs),
             "pretrained_model_name_or_path": model_lower,
             "requires_safety_checker": False,
-            "safety_checker": None,
-            "scheduler": schedulers[scheduler](**scheduler_kwargs),
-            "torch_dtype": TORCH_DTYPE,
             "use_safetensors": True,
+            "safety_checker": None,
+            "torch_dtype": dtype,
         }
 
         # already loaded
@@ -92,11 +120,19 @@ class Loader:
 
             if same_model:
                 if not same_scheduler:
-                    print(f"Swapping scheduler to {scheduler}...")
-                elif not same_karras:
+                    print(f"Switching to {scheduler}...")
+                if not same_karras:
                     print(f"{'Enabling' if karras else 'Disabling'} Karras sigmas...")
-                elif not (same_scheduler and same_karras):
+                if not same_scheduler or not same_karras:
                     self.pipe.scheduler = schedulers[scheduler](**scheduler_kwargs)
+
+                # if compiled will be an OptimizedModule
+                vae_type = type(self.pipe.vae)
+                if (issubclass(vae_type, (AutoencoderKL, OptimizedModule)) and taesd) or (
+                    issubclass(vae_type, AutoencoderTiny) and not taesd
+                ):
+                    print(f"Switching to {'Tiny' if taesd else 'KL'} VAE...")
+                    self.pipe.vae = self._load_vae(model_lower, taesd, dtype)
                 return self.pipe
             else:
                 print(f"Unloading {model_name.lower()}...")
@@ -111,39 +147,51 @@ class Loader:
         ]:
             pipe_kwargs["variant"] = "fp16"
 
-        # uses special VAE
-        if model_lower not in ["linaqruf/anything-v3-1"]:
-            pipe_kwargs["vae"] = AutoencoderTiny.from_pretrained(
-                "madebyollin/taesd",
-                torch_dtype=TORCH_DTYPE,
-                use_safetensors=True,
-            )
-
-        print(f"Loading {model_lower}...")
+        print(f"Loading {model_lower} with {'Tiny' if taesd else 'KL'} VAE...")
         self.pipe = StableDiffusionPipeline.from_pretrained(**pipe_kwargs).to(self.gpu)
+        self.pipe.vae = self._load_vae(model_lower, taesd, dtype)
+        self.pipe.load_textual_inversion(
+            pretrained_model_name_or_path=list(EMBEDDINGS.keys()),
+            tokens=list(EMBEDDINGS.values()),
+        )
         return self.pipe
 
 
-# prepare prompts for Compel
-def join_prompt(prompt: str) -> str:
-    lines = prompt.strip().splitlines()
-    return '("' + '", "'.join(lines) + '").and()' if len(lines) > 1 else prompt
+@contextmanager
+def deep_cache(pipe, interval=1, branch=0, tgate_step=0):
+    if interval > 1:
+        helper = DeepCacheSDHelper(pipe=pipe)
+        helper.set_params(cache_interval=interval, cache_branch_id=branch)
+        helper.enable()
+
+        if tgate_step > 0:
+            pipe.deepcache = helper
+            pipe.tgate = MethodType(tgate_sd_deepcache, pipe)
+
+        try:
+            yield helper
+        finally:
+            helper.disable()
+    elif interval < 2 and tgate_step > 0:
+        pipe.tgate = MethodType(tgate_sd, pipe)
+        yield None
+    else:
+        yield None
 
 
 # parse prompts with arrays
 def parse_prompt(prompt: str) -> list[str]:
-    joined_prompt = join_prompt(prompt)
-    arrays = re.findall(r"\[\[(.*?)\]\]", joined_prompt)
+    arrays = re.findall(r"\[\[(.*?)\]\]", prompt)
 
     if not arrays:
-        return [joined_prompt]
+        return [prompt]
 
     tokens = [item.split(",") for item in arrays]
     combinations = list(product(*tokens))
     prompts = []
 
     for combo in combinations:
-        current_prompt = joined_prompt
+        current_prompt = prompt
         for i, token in enumerate(combo):
             current_prompt = current_prompt.replace(f"[[{arrays[i]}]]", token.strip(), 1)
 
@@ -156,55 +204,65 @@ def generate(
     positive_prompt,
     negative_prompt="",
     seed=None,
-    model="lykon/dreamshaper-8",
+    model="Lykon/dreamshaper-8",
     scheduler="DEIS 2M",
-    aspect_ratio="1:1",
+    width=512,
+    height=512,
     guidance_scale=7.5,
     inference_steps=30,
-    karras=True,
     num_images=1,
+    karras=True,
+    taesd=False,
+    clip_skip=False,
+    truncate_prompts=False,
     increment_seed=True,
+    deep_cache_interval=1,
+    deep_cache_branch=0,
+    tgate_step=0,
     Error=Exception,
 ):
     if not torch.cuda.is_available():
         raise Error("CUDA not available")
 
-    # image dimensions
-    aspect_ratios = {
-        "16:9": (640, 360),
-        "4:3": (576, 432),
-        "1:1": (512, 512),
-        "3:4": (432, 576),
-        "9:16": (360, 640),
-    }
-    width, height = aspect_ratios[aspect_ratio]
+    if seed is None:
+        seed = int(datetime.now().timestamp())
+
+    TORCH_DTYPE = (
+        torch.bfloat16
+        if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+        else torch.float16
+    )
+
+    EMBEDDINGS_TYPE = (
+        ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NORMALIZED
+        if clip_skip
+        else ReturnedEmbeddingsType.LAST_HIDDEN_STATES_NORMALIZED
+    )
 
     with torch.inference_mode():
         loader = Loader()
-        pipe = loader.load(model, scheduler, karras)
+        pipe = loader.load(model, scheduler, karras, taesd, dtype=TORCH_DTYPE)
 
         # prompt embeds
         compel = Compel(
-            tokenizer=pipe.tokenizer,
+            textual_inversion_manager=DiffusersTextualInversionManager(pipe),
+            dtype_for_device_getter=lambda _: TORCH_DTYPE,
+            returned_embeddings_type=EMBEDDINGS_TYPE,
+            truncate_long_prompts=truncate_prompts,
             text_encoder=pipe.text_encoder,
-            truncate_long_prompts=False,
+            tokenizer=pipe.tokenizer,
             device=pipe.device,
-            dtype_for_device_getter=lambda _: TORCH_DTYPE,
         )
 
-        neg_prompt = join_prompt(negative_prompt)
-        neg_embeds = compel(neg_prompt)
-
-        if seed is None:
-            seed = int(datetime.now().timestamp())
-
-        current_seed = seed
         images = []
+        current_seed = seed
+        neg_embeds = compel(negative_prompt)
 
         for i in range(num_images):
+            # seeded generator for each iteration
             generator = torch.Generator(device=pipe.device).manual_seed(current_seed)
 
-            # run the prompt for this iteration
+            # get the prompt for this iteration
             all_positive_prompts = parse_prompt(positive_prompt)
             prompt_index = i % len(all_positive_prompts)
             pos_prompt = all_positive_prompts[prompt_index]
@@ -213,16 +271,27 @@ def generate(
                 [pos_embeds, neg_embeds]
             )
 
-            result = pipe(
-                width=width,
-                height=height,
-                prompt_embeds=pos_embeds,
-                negative_prompt_embeds=neg_embeds,
-                num_inference_steps=inference_steps,
-                guidance_scale=guidance_scale,
-                generator=generator,
-            )
-            images.append((result.images[0], str(current_seed)))
+            with deep_cache(
+                pipe,
+                interval=deep_cache_interval,
+                branch=deep_cache_branch,
+                tgate_step=tgate_step,
+            ):
+                pipe_kwargs = {
+                    "num_inference_steps": inference_steps,
+                    "negative_prompt_embeds": neg_embeds,
+                    "guidance_scale": guidance_scale,
+                    "prompt_embeds": pos_embeds,
+                    "generator": generator,
+                    "height": height,
+                    "width": width,
+                }
+                result = (
+                    pipe.tgate(**pipe_kwargs, gate_step=tgate_step)
+                    if tgate_step > 0
+                    else pipe(**pipe_kwargs)
+                )
+                images.append((result.images[0], str(current_seed)))
 
             if increment_seed:
                 current_seed += 1

requirements.txt

@@ -1,11 +1,13 @@
 accelerate
 compel
+deepcache
 diffusers
 hf-transfer
 gradio
 ruff
 scipy  # for LMS scheduler
 spaces
+tgate
 torch
 torchvision
 transformers