Add progress_bar context manager

lib/__init__.py

@@ -1,7 +1,7 @@
 from .config import Config
 from .inference import generate
 from .loader import Loader
-from .logger import Logger, log_fn
+from .logger import Logger
 from .upscaler import RealESRGAN
 from .utils import (
     async_call,
@@ -10,6 +10,7 @@ from .utils import (
     download_repo_files,
     enable_progress_bars,
     load_json,
+    progress_bar,
     read_file,
     timer,
 )
@@ -26,7 +27,7 @@ __all__ = [
     "enable_progress_bars",
     "generate",
     "load_json",
-    "log_fn",
+    "progress_bar",
     "read_file",
     "timer",
 ]

lib/inference.py

@@ -16,7 +16,7 @@ from spaces import GPU
 from .config import Config
 from .loader import Loader
 from .logger import Logger
-from .utils import load_json
+from .utils import load_json, progress_bar, timer
 
 
 def parse_prompt_with_arrays(prompt: str) -> list[str]:
@@ -193,20 +193,26 @@ def generate(
     weights = []
     loras_and_weights = [(lora_1, lora_1_weight), (lora_2, lora_2_weight)]
     loras_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "loras"))
-    for lora, weight in loras_and_weights:
-        if lora and lora.lower() != "none" and lora not in loras:
-            config = Config.CIVIT_LORAS.get(lora)
-            if config:
-                try:
-                    pipe.load_lora_weights(
-                        loras_dir,
-                        adapter_name=lora,
-                        weight_name=f"{lora}.{config['model_version_id']}.safetensors",
-                    )
-                    weights.append(weight)
-                    loras.append(lora)
-                except Exception:
-                    raise Error(f"Error loading {config['name']} LoRA")
+    total_loras = sum(1 for lora, _ in loras_and_weights if lora and lora.lower() != "none")
+    desc_loras = "Loading LoRAs"
+    if total_loras > 0:
+        with timer(f"Loading {total_loras} LoRA{'s' if total_loras > 1 else ''}"):
+            progress((0, total_loras), desc=desc_loras)
+            for i, (lora, weight) in enumerate(loras_and_weights):
+                if lora and lora.lower() != "none" and lora not in loras:
+                    config = Config.CIVIT_LORAS.get(lora)
+                    if config:
+                        try:
+                            pipe.load_lora_weights(
+                                loras_dir,
+                                adapter_name=lora,
+                                weight_name=f"{lora}.{config['model_version_id']}.safetensors",
+                            )
+                            weights.append(weight)
+                            loras.append(lora)
+                            progress((i + 1, total_loras), desc=desc_loras)
+                        except Exception:
+                            raise Error(f"Error loading {config['name']} LoRA")
 
     # unload after generating or if there was an error
     try:
@@ -294,7 +300,9 @@ def generate(
         try:
             image = pipe(**kwargs).images[0]
             if scale > 1:
-                image = upscaler.predict(image)
+                msg = f"Upscaling {scale}x"
+                with timer(msg, logger=log.info), progress_bar(100, desc=msg, progress=progress):
+                    image = upscaler.predict(image)
             images.append((image, str(current_seed)))
             current_seed += 1
         except Exception as e:

lib/loader.py

@@ -9,7 +9,7 @@ from diffusers.models.attention_processor import AttnProcessor2_0, IPAdapterAttn
 from .config import Config
 from .logger import Logger
 from .upscaler import RealESRGAN
-from .utils import timer
+from .utils import progress_bar, timer
 
 
 class Loader:
@@ -27,6 +27,20 @@ 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
+
     def _should_unload_upscaler(self, scale=1):
         if self.upscaler is not None and self.upscaler.scale != scale:
             return True
@@ -119,12 +133,15 @@ class Loader:
             setattr(self, component, None)
             gc.collect()
 
-    def _load_upscaler(self, scale=1):
+    def _load_upscaler(self, scale=1, progress=None):
         if self.upscaler is None and scale > 1:
             try:
-                with timer(f"Loading {scale}x upscaler", logger=self.log.info):
+                msg = f"Loading {scale}x upscaler"
+                # fmt: off
+                with timer(msg, logger=self.log.info), progress_bar(100, desc=msg, progress=progress):
                     self.upscaler = RealESRGAN(scale, device=self.pipe.device)
                     self.upscaler.load_weights()
+                # fmt: on
             except Exception as e:
                 self.log.error(f"Error loading {scale}x upscaler: {e}")
                 self.upscaler = None
@@ -152,9 +169,10 @@ class Loader:
             self.log.info("Enabling FreeU")
             self.pipe.enable_freeu(b1=1.5, b2=1.6, s1=0.9, s2=0.2)
 
-    def _load_ip_adapter(self, ip_adapter=""):
+    def _load_ip_adapter(self, ip_adapter="", progress=None):
         if not self.ip_adapter and ip_adapter:
-            with timer("Loading IP-Adapter", logger=self.log.info):
+            msg = "Loading IP-Adapter"
+            with timer(msg, logger=self.log.info), progress_bar(100, desc=msg, progress=progress):
                 self.pipe.load_ip_adapter(
                     "h94/IP-Adapter",
                     subfolder="models",
@@ -194,22 +212,20 @@ 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=""):
-        vae_type = type(self.pipe.vae)
-        is_kl = issubclass(vae_type, AutoencoderKL)
-        is_tiny = issubclass(vae_type, AutoencoderTiny)
-
+    def _load_vae(self, taesd=False, model="", progress=None):
         # by default all models use KL
-        if is_kl and taesd:
-            with timer("Loading Tiny VAE", logger=self.log.info):
+        if self._is_kl_vae and taesd:
+            msg = "Loading Tiny VAE"
+            with timer(msg, logger=self.log.info), progress_bar(100, desc=msg, progress=progress):
                 self.pipe.vae = AutoencoderTiny.from_pretrained(
                     pretrained_model_name_or_path="madebyollin/taesd",
                     torch_dtype=self.pipe.dtype,
                 ).to(self.pipe.device)
             return
 
-        if is_tiny and not taesd:
-            with timer("Loading KL VAE", logger=self.log.info):
+        if self._is_tiny_vae and not taesd:
+            msg = "Loading KL VAE"
+            with timer(msg, logger=self.log.info), progress_bar(100, desc=msg, progress=progress):
                 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()]}",
@@ -305,8 +321,8 @@ class Loader:
             if not same_scheduler or not same_karras:
                 self.pipe.scheduler = Config.SCHEDULERS[scheduler](**scheduler_kwargs)
 
-        self._load_vae(taesd, model)
+        self._load_vae(taesd, model, progress)
         self._load_freeu(freeu)
         self._load_deepcache(deepcache)
-        self._load_ip_adapter(ip_adapter)
-        self._load_upscaler(scale)
+        self._load_ip_adapter(ip_adapter, progress)
+        self._load_upscaler(scale, progress)

lib/logger.py

@@ -1,5 +1,4 @@
 import logging
-from functools import wraps
 from threading import Lock
 
 
@@ -54,20 +53,3 @@ class Logger:
 
     def critical(self, message, **kwargs):
         self._log(logging.CRITICAL, message, **kwargs)
-
-
-# decorator for logging function calls
-def log_fn(name=None):
-    def decorator(fn):
-        @wraps(fn)
-        def wrapper(*args, **kwargs):
-            log = Logger(name or fn.__name__)
-            log.info("begin")
-            result = fn(*args, **kwargs)
-            log.info("end")
-
-            return result
-
-        return wrapper
-
-    return decorator

lib/utils.py

@@ -35,6 +35,17 @@ def timer(message="Operation", logger=print):
         logger(f"{message} took {end - start:.2f}s")
 
 
+@contextmanager
+def progress_bar(total, desc="Loading", progress=None):
+    if progress is None:
+        yield
+    try:
+        progress((0, total), desc=desc)
+        yield
+    finally:
+        progress((total, total), desc=desc)
+
+
 @functools.lru_cache()
 def load_json(path: str) -> dict:
     with open(path, "r", encoding="utf-8") as file: