Add CLI

README.md

@@ -56,8 +56,11 @@ python -m venv .venv
 source .venv/bin/activate
 pip install -r requirements.txt torch==2.4.0 torchvision==0.19.0
 
-# http://localhost:7860
-python app.py
+# gradio
+python app.py --port 7860
+
+# cli
+python cli.py 'an astronaut riding a horse on mars'
 ```
 
 ## TODO

app.py

@@ -1,3 +1,5 @@
+import argparse
+
 import gradio as gr
 
 from generate import generate
@@ -40,7 +42,7 @@ def generate_btn_click(*args):
         prompt = None
     if prompt is None or prompt.strip() == "":
         raise gr.Error("You must enter a prompt")
-    return generate(*args)
+    return generate(*args, log=gr.Info, Error=gr.Error)
 
 
 with gr.Blocks(
@@ -87,10 +89,10 @@ with gr.Blocks(
 
                     with gr.Row():
                         num_images = gr.Dropdown(
-                            choices=list(range(1, 9)),
+                            choices=list(range(1, 5)),
                             filterable=False,
                             label="Images",
-                            value=1,
+                            value=4,
                             scale=1,
                         )
                         width = gr.Slider(
@@ -129,6 +131,7 @@ with gr.Blocks(
                     with gr.Row():
                         model = gr.Dropdown(
                             value="Lykon/dreamshaper-8",
+                            filterable=False,
                             min_width=200,
                             label="Model",
                             scale=2,
@@ -144,6 +147,7 @@ with gr.Blocks(
                         scheduler = gr.Dropdown(
                             elem_id="scheduler",
                             label="Scheduler",
+                            filterable=False,
                             value="DEIS 2M",
                             min_width=200,
                             scale=2,
@@ -186,14 +190,22 @@ with gr.Blocks(
                         tgate_step = gr.Slider(
                             label="T-GATE Step",
                             minimum=0,
-                            maximum=50,
-                            value=20,
+                            maximum=30,
+                            value=0,
                             step=1,
                         )
+
+                    with gr.Row():
+                        file_format = gr.Dropdown(
+                            choices=["png", "jpeg", "webp"],
+                            label="File Format",
+                            filterable=False,
+                            value="png",
+                        )
                         tome_ratio = gr.Slider(
                             label="ToMe Ratio",
                             minimum=0.0,
-                            maximum=1.0,
+                            maximum=0.5,
                             value=0.0,
                             step=0.01,
                         )
@@ -263,7 +275,18 @@ with gr.Blocks(
     # update the random seed using JavaScript
     random_btn.click(None, outputs=[seed], js=SEED_JS)
 
-    # ensure correct argument order
+    file_format.change(
+        lambda f: gr.Gallery(format=f),
+        inputs=[file_format],
+        outputs=[output_images],
+    )
+
+    inference_steps.change(
+        lambda max, step: gr.Slider(maximum=max, value=min(max, step)),
+        inputs=[inference_steps, tgate_step],
+        outputs=[tgate_step],
+    )
+
     generate_btn.click(
         generate_btn_click,
         api_name="api",
@@ -291,8 +314,14 @@ with gr.Blocks(
         ],
     )
 
-# https://www.gradio.app/docs/gradio/interface#interface-queue
-demo.queue().launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
+    parser.add_argument("-s", "--server", type=str, metavar="STR", default="0.0.0.0")
+    parser.add_argument("-p", "--port", type=int, metavar="INT", default=7860)
+    args = parser.parse_args()
+
+    # https://www.gradio.app/docs/gradio/interface#interface-queue
+    demo.queue().launch(
+        server_name=args.server,
+        server_port=args.port,
+    )

cli.py

@@ -0,0 +1,59 @@
+import argparse
+
+from generate import generate
+
+
+def save_images(images, filename="image.png"):
+    for i, (img, _) in enumerate(images):
+        name, ext = filename.rsplit(".", 1)
+        img.save(f"{name}.{ext}" if len(images) == 1 else f"{name}_{i}.{ext}")
+
+
+def main():
+    parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
+    parser.add_argument("prompt", type=str, metavar="PROMPT")
+    parser.add_argument("-n", "--negative", type=str, metavar="STR", default="<fast_negative>")
+    parser.add_argument("-s", "--seed", type=int, metavar="INT")
+    parser.add_argument("-i", "--images", type=int, metavar="INT", default=1)
+    parser.add_argument("-f", "--filename", type=str, metavar="STR", default="image.png")
+    parser.add_argument("-w", "--width", type=int, metavar="INT", default=448)
+    parser.add_argument("-h", "--height", type=int, metavar="INT", default=576)
+    parser.add_argument("-m", "--model", type=str, metavar="STR", default="Lykon/dreamshaper-8")
+    parser.add_argument("-d", "--deepcache", type=int, metavar="INT", default=2)
+    parser.add_argument("-t", "--tgate", type=int, metavar="INT", default=20)
+    parser.add_argument("--scheduler", type=str, metavar="STR", default="DEIS 2M")
+    parser.add_argument("--guidance", type=float, metavar="FLOAT", default=7)
+    parser.add_argument("--steps", type=int, metavar="INT", default=30)
+    parser.add_argument("--tome", type=float, metavar="FLOAT", default=0.0)
+    parser.add_argument("--taesd", action="store_true")
+    parser.add_argument("--clip-skip", action="store_true")
+    parser.add_argument("--truncate", action="store_true")
+    parser.add_argument("--no-karras", action="store_false")
+    parser.add_argument("--no-increment", action="store_false")
+
+    args = parser.parse_args()
+    images = generate(
+        args.prompt,
+        args.negative,
+        args.seed,
+        args.model,
+        args.scheduler,
+        args.width,
+        args.height,
+        args.guidance,
+        args.steps,
+        args.images,
+        args.no_karras,
+        args.taesd,
+        args.clip_skip,
+        args.truncate,
+        args.no_increment,
+        args.deepcache,
+        args.tgate,
+        args.tome,
+    )
+    save_images(images, args.filename)
+
+
+if __name__ == "__main__":
+    main()

generate.py

@@ -5,9 +5,9 @@ from datetime import datetime
 from itertools import product
 from os import environ
 from types import MethodType
+from typing import Callable
 from warnings import filterwarnings
 
-import gradio as gr
 import spaces
 import tomesd
 import torch
@@ -80,7 +80,6 @@ class Loader:
                 tgate_sd_deepcache if has_deepcache else tgate_sd,
                 self.pipe,
             )
-
         return self.pipe.tgate
 
     def _load_vae(self, model_name=None, taesd=False, dtype=None):
@@ -244,10 +243,11 @@ def generate(
     deepcache_interval=1,
     tgate_step=0,
     tome_ratio=0,
-    progress=gr.Progress(track_tqdm=True),
+    log: Callable[[str], None] = None,
+    Error=Exception,
 ):
     if not torch.cuda.is_available():
-        raise gr.Error("CUDA not available")
+        raise Error("CUDA not available")
 
     if seed is None:
         seed = int(datetime.now().timestamp())
@@ -324,5 +324,6 @@ def generate(
 
         end = time.perf_counter()
         diff = end - start
-        gr.Info(f"Generated {len(images)} image{'s' if len(images) > 1 else ''} in {diff:.2f}s")
+        if log:
+            log(f"Generated {len(images)} image{'s' if len(images) > 1 else ''} in {diff:.2f}s")
         return images

usage.md

@@ -41,7 +41,7 @@ When using arrays, you should disable `Autoincrement` so the same seed is used f
 
 #### Schedulers
 
-All are based on [k_diffusion](https://github.com/crowsonkb/k-diffusion) except [DEIS](https://github.com/qsh-zh/deis) and [DPM++](https://github.com/LuChengTHU/dpm-solver). Optionally, the [Karras](https://arxiv.org/abs/2206.00364) noise schedule can be used:
+Optionally, the [Karras](https://arxiv.org/abs/2206.00364) noise schedule can be used:
 
 * [DEIS 2M](https://huggingface.co/docs/diffusers/en/api/schedulers/deis) (default)
 * [DPM++ 2M](https://huggingface.co/docs/diffusers/en/api/schedulers/multistep_dpm_solver)
@@ -63,11 +63,11 @@ All are based on [k_diffusion](https://github.com/crowsonkb/k-diffusion) except
 
 #### T-GATE
 
-[T-GATE](https://github.com/HaozheLiu-ST/T-GATE) (Zhang et al. 2024) caches self and cross attention computations up to `Step`. Afterwards, attention is no longer computed and the cache is used, resulting in a noticeable speedup. Defaults to `20`.
+[Temporal gating](https://github.com/HaozheLiu-ST/T-GATE) (Zhang et al. 2024) caches self and cross attention computations up to `Step`. Afterwards, attention is no longer computed and the cache is used, resulting in a noticeable speedup.
 
-#### ToME
+#### ToMe
 
-[ToMe](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://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.
 
 #### Tiny VAE
 
@@ -79,4 +79,4 @@ When enabled, the last CLIP layer is skipped. This _can_ improve image quality w
 
 #### 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.
+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.