From 7ae77faf07db11acc0cd978924f4dec23e7a4736 Mon Sep 17 00:00:00 2001
From: Bryce <github20210803@accounts.brycedrennan.com>
Date: Mon, 10 Oct 2022 20:13:32 -0700
Subject: [PATCH] refactor: log indentation via formatter

Replace manual indentation with a context-aware logging formatter.
---
 imaginairy/api.py                                   | 12 ++++++------
 imaginairy/enhancers/face_restoration_codeformer.py |  2 +-
 imaginairy/log_utils.py                             | 13 ++++++++++++-
 imaginairy/schema.py                                |  2 +-
 4 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/imaginairy/api.py b/imaginairy/api.py
index 8ac6509..704cb53 100755
--- a/imaginairy/api.py
+++ b/imaginairy/api.py
@@ -147,7 +147,7 @@ def imagine_image_files(
                 subpath, f"{basefilename}_[{image_type}].{output_file_extension}"
             )
             result.save(filepath, image_type=image_type)
-            logger.info(f"    🖼  [{image_type}] saved to: {filepath}")
+            logger.info(f"🖼  [{image_type}] saved to: {filepath}")
         base_count += 1
         del result
 
@@ -181,12 +181,12 @@ def imagine(
         precision
     ), fix_torch_nn_layer_norm(), fix_torch_group_norm():
         for prompt in prompts:
+            logger.info(f"Generating 🖼  : {prompt.prompt_description()}")
             with ImageLoggingContext(
                 prompt=prompt,
                 model=model,
                 img_callback=img_callback,
             ):
-                logger.info(f"Generating {prompt.prompt_description()}")
                 seed_everything(prompt.seed)
                 model.tile_mode(prompt.tile_mode)
 
@@ -234,7 +234,7 @@ def imagine(
                             max_width=prompt.width,
                         )
                     except PIL.UnidentifiedImageError:
-                        logger.warning(f"   Could not load image: {prompt.init_image}")
+                        logger.warning(f"Could not load image: {prompt.init_image}")
                         continue
 
                     init_image_t = pillow_img_to_torch_image(init_image)
@@ -352,7 +352,7 @@ def imagine(
 
                     if add_caption:
                         caption = generate_caption(img)
-                        logger.info(f"    Generated caption: {caption}")
+                        logger.info(f"Generated caption: {caption}")
 
                     safety_score = create_safety_score(
                         img,
@@ -360,10 +360,10 @@ def imagine(
                     )
                     if not safety_score.is_filtered:
                         if prompt.fix_faces:
-                            logger.info("    Fixing 😊 's in 🖼  using CodeFormer...")
+                            logger.info("Fixing 😊 's in 🖼  using CodeFormer...")
                             img = enhance_faces(img, fidelity=prompt.fix_faces_fidelity)
                         if prompt.upscale:
-                            logger.info("    Upscaling 🖼  using real-ESRGAN...")
+                            logger.info("Upscaling 🖼  using real-ESRGAN...")
                             upscaled_img = upscale_image(img)
 
                         # put the newly generated patch back into the original, full size image
diff --git a/imaginairy/enhancers/face_restoration_codeformer.py b/imaginairy/enhancers/face_restoration_codeformer.py
index 9db766a..145343e 100644
--- a/imaginairy/enhancers/face_restoration_codeformer.py
+++ b/imaginairy/enhancers/face_restoration_codeformer.py
@@ -67,7 +67,7 @@ def enhance_faces(img, fidelity=0):
     num_det_faces = face_helper.get_face_landmarks_5(
         only_center_face=False, resize=640, eye_dist_threshold=5
     )
-    logger.info(f"    Enhancing {num_det_faces} faces")
+    logger.info(f"Enhancing {num_det_faces} faces")
     # align and warp each face
     face_helper.align_warp_face()
 
diff --git a/imaginairy/log_utils.py b/imaginairy/log_utils.py
index 2523daa..01d1850 100644
--- a/imaginairy/log_utils.py
+++ b/imaginairy/log_utils.py
@@ -98,6 +98,14 @@ def conditioning_to_img(conditioning):
     return ToPILImage()(conditioning)
 
 
+class IndentingFormatter(logging.Formatter):
+    def format(self, record):
+        s = super().format(record)
+        if _CURRENT_LOGGING_CONTEXT is not None:
+            s = f"    {s}"
+        return s
+
+
 def configure_logging(level="INFO"):
     fmt = "%(message)s"
     if level == "DEBUG":
@@ -107,7 +115,10 @@ def configure_logging(level="INFO"):
         "version": 1,
         "disable_existing_loggers": True,
         "formatters": {
-            "standard": {"format": fmt},
+            "standard": {
+                "format": fmt,
+                "class": "imaginairy.log_utils.IndentingFormatter",
+            },
         },
         "handlers": {
             "default": {
diff --git a/imaginairy/schema.py b/imaginairy/schema.py
index ff103f6..82a6be4 100644
--- a/imaginairy/schema.py
+++ b/imaginairy/schema.py
@@ -153,7 +153,7 @@ class ImaginePrompt:
 
     def prompt_description(self):
         return (
-            f'🖼  : "{self.prompt_text}" {self.width}x{self.height}px '
+            f'"{self.prompt_text}" {self.width}x{self.height}px '
             f"seed:{self.seed} prompt-strength:{self.prompt_strength} steps:{self.steps} sampler-type:{self.sampler_type}"
         )