Better logging

.gitignore

@@ -1,3 +1,4 @@
 __pycache__/
 .venv/
 loras/
+app.log

app.py

@@ -2,11 +2,19 @@ import argparse
 import json
 import os
 import random
+from warnings import filterwarnings
 
 import gradio as gr
+from diffusers.utils import logging as diffusers_logging
+from transformers import logging as transformers_logging
 
 from lib import Config, async_call, download_civit_file, download_repo_files, generate, read_file
 
+filterwarnings("ignore", category=FutureWarning, module="diffusers")
+filterwarnings("ignore", category=FutureWarning, module="transformers")
+diffusers_logging.set_verbosity_error()
+transformers_logging.set_verbosity_error()
+
 # the CSS `content` attribute expects a string so we need to wrap the number in quotes
 refresh_seed_js = """
 () => {

data/prompts.json

@@ -32,7 +32,7 @@
   "bowl of steaming hot ramen with a sliced egg, thin slices of meat, green onions, noodles, chopsticks, solo, minimal",
   "large pizza with melted cheese, seared pepperoni, crispy crust, solo, minimal",
   "sizzling hot sirloin steak with a perfect crust, seared to perfection, served with a side of roasted vegetables and mashed potatoes, solo, minimal",
-  "three-tier wedding cake, white frosting, colorful accent flowers, fresh berry garnish, elegant, minimal",
+  "wedding cake, white frosting, colorful accent flowers, fresh berry garnish, tiers, layers, elegant, minimal",
   "baked salmon fillet with a perfectly crispy skin and flaky flesh, side of steamed vegetables and quinoa, healthy, fresh, solo, minimal",
   "steaming bowl of hearty chili with tender chunks of beef, rich tomato sauce, topped with grated cheddar cheese and green onions, solo, minimal",
   "platter of sushi rolls, tuna, salmon, california maki, rainbow, colorful, beautiful arrangement, solo, minimal",

lib/__init__.py

@@ -1,17 +1,20 @@
 from .config import Config
 from .inference import generate
 from .loader import Loader
+from .logger import Logger, log_fn
 from .upscaler import RealESRGAN
 from .utils import async_call, download_civit_file, download_repo_files, load_json, read_file
 
 __all__ = [
     "Config",
     "Loader",
+    "Logger",
     "RealESRGAN",
     "async_call",
     "download_civit_file",
     "download_repo_files",
     "generate",
     "load_json",
+    "log_fn",
     "read_file",
 ]

lib/inference.py

@@ -15,11 +15,9 @@ from PIL import Image
 
 from .config import Config
 from .loader import Loader
+from .logger import Logger
 from .utils import load_json
 
-__import__("warnings").filterwarnings("ignore", category=FutureWarning, module="transformers")
-__import__("transformers").logging.set_verbosity_error()
-
 
 def parse_prompt_with_arrays(prompt: str) -> list[str]:
     arrays = re.findall(r"\[\[(.*?)\]\]", prompt)
@@ -170,6 +168,7 @@ def generate(
         return latents
 
     start = time.perf_counter()
+    log = Logger("generate")
     loader = Loader()
     loader.load(
         KIND,
@@ -249,10 +248,9 @@ def generate(
     images = []
     current_seed = seed
     for i in range(num_images):
-        # seeded generator for each iteration
-        generator = torch.Generator(device=pipe.device).manual_seed(current_seed)
-
         try:
+            generator = torch.Generator(device=pipe.device).manual_seed(current_seed)
+
             positive_prompts = parse_prompt_with_arrays(positive_prompt)
             index = i % len(positive_prompts)
             positive_styled, negative_styled = apply_style(
@@ -270,7 +268,6 @@ def generate(
             for embedding in embeddings:
                 negative_styled += f", <{embedding}>"
 
-            # print prompts
             positive_embeds, negative_embeds = compel.pad_conditioning_tensors_to_same_length(
                 [compel(positive_styled), compel(negative_styled)]
             )
@@ -317,6 +314,8 @@ def generate(
             CURRENT_IMAGE += 1
 
     diff = time.perf_counter() - start
+    msg = f"Generated {len(images)} image{'s' if len(images) > 1 else ''} in {diff:.2f}s"
+    log.info(msg)
     if Info:
-        Info(f"Generated {len(images)} image{'s' if len(images) > 1 else ''} in {diff:.2f}s")
+        Info(msg)
     return images

lib/loader.py

@@ -1,21 +1,16 @@
 import gc
 from threading import Lock
-from warnings import filterwarnings
 
 import torch
 from DeepCache import DeepCacheSDHelper
 from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionPipeline
 from diffusers.models import AutoencoderKL, AutoencoderTiny
 from diffusers.models.attention_processor import AttnProcessor2_0, IPAdapterAttnProcessor2_0
-from torch._dynamo import OptimizedModule
 
 from .config import Config
+from .logger import Logger
 from .upscaler import RealESRGAN
 
-__import__("diffusers").logging.set_verbosity_error()
-filterwarnings("ignore", category=FutureWarning, module="torch")
-filterwarnings("ignore", category=FutureWarning, module="diffusers")
-
 
 class Loader:
     _instance = None
@@ -30,10 +25,16 @@ class Loader:
                 cls._instance.ip_adapter = None
                 cls._instance.upscaler_2x = None
                 cls._instance.upscaler_4x = None
+                cls._instance.log = Logger("Loader")
         return cls._instance
 
-    def _should_unload_ip_adapter(self, ip_adapter=""):
-        return self.ip_adapter is not None and not ip_adapter
+    def _should_unload_ip_adapter(self, model="", ip_adapter=""):
+        # unload if model changed
+        if self.model and self.model.lower() != model.lower():
+            return True
+        if self.ip_adapter and not ip_adapter:
+            return True
+        return False
 
     def _should_unload_pipeline(self, kind="", model=""):
         if self.pipe is None:
@@ -48,7 +49,10 @@ class Loader:
 
     # https://github.com/huggingface/diffusers/blob/v0.28.0/src/diffusers/loaders/ip_adapter.py#L300
     def _unload_ip_adapter(self):
-        print("Unloading IP Adapter...")
+        if self.ip_adapter is None:
+            return
+
+        self.log.info("Unloading IP-Adapter")
         if not isinstance(self.pipe, StableDiffusionImg2ImgPipeline):
             self.pipe.image_encoder = None
             self.pipe.register_to_config(image_encoder=[None, None])
@@ -72,13 +76,12 @@ class Loader:
         gc.collect()
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
-        torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()
         torch.cuda.synchronize()
 
     def _unload(self, kind="", model="", ip_adapter=""):
         to_unload = []
-        if self._should_unload_ip_adapter(ip_adapter):
+        if self._should_unload_ip_adapter(model, ip_adapter):
             self._unload_ip_adapter()
             to_unload.append("ip_adapter")
         if self._should_unload_pipeline(kind, model):
@@ -91,8 +94,8 @@ class Loader:
             setattr(self, component, None)
 
     def _load_ip_adapter(self, ip_adapter=""):
-        if self.ip_adapter is None and ip_adapter:
-            print(f"Loading IP Adapter: {ip_adapter}...")
+        if not self.ip_adapter and ip_adapter:
+            self.log.info(f"Loading IP-Adapter: {ip_adapter}")
             self.pipe.load_ip_adapter(
                 "h94/IP-Adapter",
                 subfolder="models",
@@ -102,29 +105,30 @@ class Loader:
             self.pipe.set_ip_adapter_scale(0.5)
             self.ip_adapter = ip_adapter
 
+    # upscalers don't need to be unloaded
     def _load_upscaler(self, scale=1):
         if scale == 2 and self.upscaler_2x is None:
             try:
-                print("Loading 2x upscaler...")
+                self.log.info("Loading 2x upscaler")
                 self.upscaler_2x = RealESRGAN(2, "cuda")
                 self.upscaler_2x.load_weights()
             except Exception as e:
-                print(f"Error loading 2x upscaler: {e}")
+                self.log.error(f"Error loading 2x upscaler: {e}")
                 self.upscaler_2x = None
         if scale == 4 and self.upscaler_4x is None:
             try:
-                print("Loading 4x upscaler...")
+                self.log.info("Loading 4x upscaler")
                 self.upscaler_4x = RealESRGAN(4, "cuda")
                 self.upscaler_4x.load_weights()
             except Exception as e:
-                print(f"Error loading 4x upscaler: {e}")
+                self.log.error(f"Error loading 4x upscaler: {e}")
                 self.upscaler_4x = None
 
     def _load_pipeline(self, kind, model, tqdm, **kwargs):
         pipeline = Config.PIPELINES[kind]
         if self.pipe is None:
             try:
-                print(f"Loading {model}...")
+                self.log.info(f"Loading {model}")
                 self.model = model
                 if model.lower() in Config.MODEL_CHECKPOINTS.keys():
                     self.pipe = pipeline.from_single_file(
@@ -134,7 +138,7 @@ class Loader:
                 else:
                     self.pipe = pipeline.from_pretrained(model, **kwargs).to("cuda")
             except Exception as e:
-                print(f"Error loading {model}: {e}")
+                self.log.error(f"Error loading {model}: {e}")
                 self.model = None
                 self.pipe = None
                 return
@@ -145,38 +149,32 @@ class Loader:
 
     def _load_vae(self, taesd=False, model=""):
         vae_type = type(self.pipe.vae)
-        is_kl = issubclass(vae_type, (AutoencoderKL, OptimizedModule))
+        is_kl = issubclass(vae_type, AutoencoderKL)
         is_tiny = issubclass(vae_type, AutoencoderTiny)
 
         # by default all models use KL
         if is_kl and taesd:
-            print("Switching to Tiny VAE...")
+            self.log.info("Switching to Tiny VAE")
             self.pipe.vae = AutoencoderTiny.from_pretrained(
-                # can't compile tiny VAE
                 pretrained_model_name_or_path="madebyollin/taesd",
                 torch_dtype=self.pipe.dtype,
             ).to(self.pipe.device)
             return
 
         if is_tiny and not taesd:
-            print("Switching to KL VAE...")
+            self.log.info("Switching to KL VAE")
             if model.lower() in Config.MODEL_CHECKPOINTS.keys():
-                vae = AutoencoderKL.from_single_file(
+                self.pipe.vae = AutoencoderKL.from_single_file(
                     f"https://huggingface.co/{model}/{Config.MODEL_CHECKPOINTS[model.lower()]}",
                     torch_dtype=self.pipe.dtype,
                 ).to(self.pipe.device)
             else:
-                vae = AutoencoderKL.from_pretrained(
+                self.pipe.vae = AutoencoderKL.from_pretrained(
                     pretrained_model_name_or_path=model,
                     torch_dtype=self.pipe.dtype,
                     subfolder="vae",
                     variant="fp16",
                 ).to(self.pipe.device)
-            self.pipe.vae = torch.compile(
-                mode="reduce-overhead",
-                fullgraph=True,
-                model=vae,
-            )
 
     def _load_deepcache(self, interval=1):
         has_deepcache = hasattr(self.pipe, "deepcache")
@@ -185,6 +183,7 @@ class Loader:
         if has_deepcache:
             self.pipe.deepcache.disable()
         else:
+            self.log.info("Loading DeepCache")
             self.pipe.deepcache = DeepCacheSDHelper(pipe=self.pipe)
         self.pipe.deepcache.set_params(cache_interval=interval)
         self.pipe.deepcache.enable()
@@ -195,10 +194,10 @@ class Loader:
         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:
-            print("Disabling FreeU...")
+            self.log.info("Disabling FreeU")
             self.pipe.disable_freeu()
         elif not has_freeu and freeu:
-            print("Enabling FreeU...")
+            self.log.info("Enabling FreeU")
             self.pipe.enable_freeu(b1=1.5, b2=1.6, s1=0.9, s2=0.2)
 
     def load(
@@ -271,14 +270,14 @@ class Loader:
         # same model, different scheduler
         if self.model.lower() == model.lower():
             if not same_scheduler:
-                print(f"Switching to {scheduler}...")
+                self.log.info(f"Switching to {scheduler}")
             if not same_karras:
-                print(f"{'Enabling' if karras else 'Disabling'} Karras sigmas...")
+                self.log.info(f"{'Enabling' if karras else 'Disabling'} Karras sigmas")
             if not same_scheduler or not same_karras:
                 self.pipe.scheduler = Config.SCHEDULERS[scheduler](**scheduler_kwargs)
 
-        self._load_freeu(freeu)
         self._load_vae(taesd, model)
+        self._load_upscaler(scale)
+        self._load_freeu(freeu)
         self._load_deepcache(deepcache)
         self._load_ip_adapter(ip_adapter)
-        self._load_upscaler(scale)

lib/logger.py

@@ -0,0 +1,73 @@
+import logging
+from functools import wraps
+from threading import Lock
+
+
+class Logger:
+    _instances = {}
+    _lock = Lock()
+
+    def __new__(cls, name="root"):
+        with cls._lock:
+            if name not in cls._instances:
+                instance = super().__new__(cls)
+                instance._init(name)
+                cls._instances[name] = instance
+            return cls._instances[name]
+
+    def _init(self, name):
+        self.logger = logging.getLogger(name)
+        self.logger.setLevel(logging.DEBUG)
+        self.logger.propagate = False
+
+        console_handler = logging.StreamHandler()
+        console_handler.setLevel(logging.INFO)
+
+        file_handler = logging.FileHandler("app.log")
+        file_handler.setLevel(logging.DEBUG)
+
+        formatter = logging.Formatter(
+            "%(asctime)s [%(threadName)s] %(levelname)-5s %(name)s - %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",  # no milliseconds
+        )
+        console_handler.setFormatter(formatter)
+        file_handler.setFormatter(formatter)
+
+        self.logger.addHandler(console_handler)
+        self.logger.addHandler(file_handler)
+
+    def _log(self, level, message):
+        log_message = f"{message}".strip()
+        self.logger.log(level, log_message)
+
+    def debug(self, message, **kwargs):
+        self._log(logging.DEBUG, message, **kwargs)
+
+    def info(self, message, **kwargs):
+        self._log(logging.INFO, message, **kwargs)
+
+    def warning(self, message, **kwargs):
+        self._log(logging.WARNING, message, **kwargs)
+
+    def error(self, message, **kwargs):
+        self._log(logging.ERROR, message, **kwargs)
+
+    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