Async generate wrapper

app.py

@@ -4,7 +4,7 @@ import random
 
 import gradio as gr
 
-from lib import Config, generate
+from lib import Config, async_call, generate
 
 # the CSS `content` attribute expects a string so we need to wrap the number in quotes
 refresh_seed_js = """
@@ -79,7 +79,7 @@ def image_select_fn(images, image, i):
     return gr.Image(images[i][0]) if i > -1 else None
 
 
-def generate_fn(*args):
+async def generate_fn(*args):
     if len(args) > 0:
         prompt = args[0]
     else:
@@ -87,7 +87,7 @@ def generate_fn(*args):
     if prompt is None or prompt.strip() == "":
         raise gr.Error("You must enter a prompt")
     try:
-        images = generate(*args, Info=gr.Info, Error=gr.Error)
+        images = await async_call(generate, *args, Info=gr.Info, Error=gr.Error)
     except RuntimeError:
         raise gr.Error("RuntimeError: Please try again")
     return images
@@ -194,25 +194,25 @@ with gr.Blocks(
                         width = gr.Slider(
                             value=Config.WIDTH,
                             label="Width",
-                            minimum=320,
+                            minimum=256,
                             maximum=768,
-                            step=16,
+                            step=32,
                         )
                         height = gr.Slider(
                             value=Config.HEIGHT,
                             label="Height",
-                            minimum=320,
+                            minimum=256,
                             maximum=768,
-                            step=16,
+                            step=32,
                         )
                         aspect_ratio = gr.Dropdown(
                             choices=[
                                 ("Custom", None),
+                                ("4:7 (384x672)", "384,672"),
                                 ("7:9 (448x576)", "448,576"),
-                                ("3:4 (432x576)", "432,576"),
                                 ("1:1 (512x512)", "512,512"),
-                                ("4:3 (576x432)", "576,432"),
                                 ("9:7 (576x448)", "576,448"),
+                                ("7:4 (672x384)", "672,384"),
                             ],
                             value="448,576",
                             filterable=False,

cli.py

@@ -1,8 +1,9 @@
 # CLI
 # usage: python cli.py 'colorful calico cat artstation'
 import argparse
+import asyncio
 
-from lib import Config, generate
+from lib import Config, async_call, generate
 
 
 def save_images(images, filename="image.png"):
@@ -11,7 +12,7 @@ def save_images(images, filename="image.png"):
         img.save(f"{name}.{ext}" if len(images) == 1 else f"{name}_{i}.{ext}")
 
 
-def main():
+async def main():
     # fmt: off
     parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
     parser.add_argument("prompt", type=str, metavar="PROMPT")
@@ -42,7 +43,8 @@ def main():
     # fmt: on
 
     args = parser.parse_args()
-    images = generate(
+    images = await async_call(
+        generate,
         args.prompt,
         args.negative,
         args.image,
@@ -68,8 +70,8 @@ def main():
         args.deepcache,
         args.scale,
     )
-    save_images(images, args.filename)
+    await async_call(save_images, images, args.filename)
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())

lib/__init__.py

@@ -1,6 +1,6 @@
 from .config import Config
-from .inference import generate
+from .inference import async_call, generate
 from .loader import Loader
 from .upscaler import RealESRGAN
 
-__all__ = ["Config", "Loader", "RealESRGAN", "generate"]
+__all__ = ["Config", "Loader", "RealESRGAN", "async_call", "generate"]

lib/config.py

@@ -41,7 +41,7 @@ Config = SimpleNamespace(
     GUIDANCE_SCALE=6,
     INFERENCE_STEPS=35,
     DENOISING_STRENGTH=0.6,
-    DEEPCACHE_INTERVAL=2,
+    DEEPCACHE_INTERVAL=1,
     SCALE=1,
     SCALES=[1, 2, 4],
 )

lib/inference.py

@@ -1,26 +1,48 @@
+import functools
+import inspect
 import json
 import os
 import re
 import time
 from datetime import datetime
 from itertools import product
-from typing import Callable
+from typing import Callable, TypeVar
 
+import anyio
 import numpy as np
 import spaces
 import torch
+from anyio import Semaphore
 from compel import Compel, DiffusersTextualInversionManager, ReturnedEmbeddingsType
 from compel.prompt_parser import PromptParser
 from huggingface_hub.utils import HFValidationError, RepositoryNotFoundError
 from PIL import Image
+from typing_extensions import ParamSpec
 
 from .loader import Loader
 
 __import__("warnings").filterwarnings("ignore", category=FutureWarning, module="transformers")
 __import__("transformers").logging.set_verbosity_error()
 
+T = TypeVar("T")
+P = ParamSpec("P")
+
+MAX_CONCURRENT_THREADS = 1
+MAX_THREADS_GUARD = Semaphore(MAX_CONCURRENT_THREADS)
+
 with open("./data/styles.json") as f:
-    styles = json.load(f)
+    STYLES = json.load(f)
+
+
+# like the original but supports args and kwargs instead of a dict
+# https://github.com/huggingface/huggingface-inference-toolkit/blob/0.2.0/src/huggingface_inference_toolkit/async_utils.py
+async def async_call(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
+    async with MAX_THREADS_GUARD:
+        sig = inspect.signature(fn)
+        bound_args = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+        partial_fn = functools.partial(fn, **bound_args.arguments)
+        return await anyio.to_thread.run_sync(partial_fn)
 
 
 # parse prompts with arrays
@@ -43,10 +65,10 @@ def parse_prompt(prompt: str) -> list[str]:
 
 
 def apply_style(prompt, style_id, negative=False):
-    global styles
+    global STYLES
     if not style_id or style_id == "None":
         return prompt
-    for style in styles:
+    for style in STYLES:
         if style["id"] == style_id:
             if negative:
                 return prompt + " . " + style["negative_prompt"]
@@ -55,7 +77,7 @@ def apply_style(prompt, style_id, negative=False):
     return prompt
 
 
-def prepare_image(input, size=(512, 512)):
+def prepare_image(input, size=None):
     image = None
     if isinstance(input, Image.Image):
         image = input
@@ -65,7 +87,11 @@ def prepare_image(input, size=(512, 512)):
         if os.path.isfile(input):
             image = Image.open(input)
     if image is not None:
-        return image.convert("RGB").resize(size, Image.Resampling.LANCZOS)
+        image = image.convert("RGB")
+    if size is not None:
+        image = image.resize(size, Image.Resampling.LANCZOS)
+    if image is not None:
+        return image
     else:
         raise ValueError("Invalid image prompt")
 
@@ -213,7 +239,9 @@ def generate(
                 kwargs["image"] = prepare_image(image_prompt, (width, height))
 
             if IP_ADAPTER:
-                kwargs["ip_adapter_image"] = prepare_image(ip_image, (width, height))
+                # don't resize full-face images
+                size = None if ip_face else (width, height)
+                kwargs["ip_adapter_image"] = prepare_image(ip_image, size)
 
             try:
                 image = pipe(**kwargs).images[0]

lib/loader.py

@@ -104,31 +104,33 @@ class Loader:
             print("Switching to Tiny VAE...")
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 pretrained_model_name_or_path="madebyollin/taesd",
-            ).to(self.pipe.device, self.pipe.dtype)
+                torch_dtype=self.pipe.dtype,
+            ).to(self.pipe.device)
             return
 
         if is_tiny and not taesd:
             print("Switching to KL VAE...")
             model = AutoencoderKL.from_pretrained(
                 pretrained_model_name_or_path=model_name,
+                torch_dtype=self.pipe.dtype,
                 subfolder="vae",
                 variant=variant,
-            ).to(self.pipe.device, self.pipe.dtype)
+            ).to(self.pipe.device)
             self.pipe.vae = torch.compile(
                 mode="reduce-overhead",
                 fullgraph=True,
                 model=model,
             )
 
-    def _load_pipeline(self, kind, model, device, dtype, **kwargs):
+    def _load_pipeline(self, kind, model, device, **kwargs):
         pipelines = {
             "txt2img": StableDiffusionPipeline,
             "img2img": StableDiffusionImg2ImgPipeline,
         }
         if self.pipe is None:
-            self.pipe = pipelines[kind].from_pretrained(model, **kwargs).to(device, dtype)
+            self.pipe = pipelines[kind].from_pretrained(model, **kwargs).to(device)
         if not isinstance(self.pipe, pipelines[kind]):
-            self.pipe = pipelines[kind].from_pipe(self.pipe).to(device, dtype)
+            self.pipe = pipelines[kind].from_pipe(self.pipe).to(device)
             self.ip_adapter = None
 
     def load(
@@ -186,13 +188,14 @@ class Loader:
             "scheduler": schedulers[scheduler](**scheduler_kwargs),
             "requires_safety_checker": False,
             "safety_checker": None,
+            "torch_dtype": dtype,
             "variant": variant,
         }
 
         if self.pipe is None:
             print(f"Loading {model_lower} with {'Tiny' if taesd else 'KL'} VAE...")
 
-        self._load_pipeline(kind, model_lower, device, dtype, **pipe_kwargs)
+        self._load_pipeline(kind, model_lower, device, **pipe_kwargs)
         model_name = self.pipe.config._name_or_path
         same_model = model_name.lower() == model_lower
         same_scheduler = isinstance(self.pipe.scheduler, schedulers[scheduler])
@@ -210,7 +213,7 @@ class Loader:
                 self.pipe.scheduler = schedulers[scheduler](**scheduler_kwargs)
         else:
             self.pipe = None
-            self._load_pipeline(kind, model_lower, device, dtype, **pipe_kwargs)
+            self._load_pipeline(kind, model_lower, device, **pipe_kwargs)
 
         self._load_ip_adapter(ip_adapter)
         self._load_vae(taesd, model_lower, variant)

requirements.txt

@@ -1,3 +1,4 @@
+anyio==4.4.0
 accelerate
 einops==0.8.0
 compel==2.0.3