Improve readability of output, reduce diffusers logging clutter.

diff --git a/test.py b/test.py
index ef64c446a98ff161c0e69a21f41b04f206b59b36..d485214de9ee9e116ea3733fbdabc83a9a91cda0 100755 (executable)
--- a/test.py
+++ b/test.py
@@ -6,9 +6,22 @@ from exiftool import ExifToolHelper
 from torch import Generator
 from diffusers import StableDiffusionPipeline
 from os.path import dirname, basename, splitext, join as path_join, exists
+from logging import Formatter as LoggingFormatter, captureWarnings, Filter as LoggingFilter
+from diffusers.utils import logging
 
 DEFAULT_MODEL='./v1-5-pruned-emaonly.safetensors'
 
+
+class FilterOutString(LoggingFilter):
+
+    def __init__(self, target):
+        super().__init__()
+        self.target = target
+
+    def filter(self, record):
+        return self.target not in record.getMessage()
+
+
 def parse_args():
     parser = ArgumentParser(add_help=False)
     parser.add_argument('-p', '--prompt', required=True)
@@ -39,17 +52,25 @@ for n in range(args.number):
     if exists(path):
         raise Exception(f'Would overwrite file: {path}')
 
-pipe = StableDiffusionPipeline.from_single_file(args.model)
+print(f'SETTING UP STABLE DIFFUSION PIPELINE FROM MODEL: {args.model}\n')
+diffusers_logging_handler = logging.get_logger('diffusers').handlers[0]
+diffusers_logging_handler.setFormatter(LoggingFormatter(fmt='DIFFUSERS WARNING: %(message)s\n'))
+diffusers_logging_handler.addFilter(FilterOutString('You have disabled the safety checker'))
+captureWarnings(True)
+logging.disable_progress_bar()
+pipe = StableDiffusionPipeline.from_single_file(args.model, local_files_only=True)
+print('PIPELINE READY\n')
+
 pipe.to('cuda')
 generator = Generator()
 start_seed = randint(-(2**31-1), 2**31) if args.randomness_seed == 0 else args.randomness_seed
 for n in range(args.number):
     seed = start_seed + n
+    path = save_path(n)
+    metadata = f'SEED: {seed}; GUIDANCE: {args.guidance}; HEIGHT: {args.height}; WIDTH: {args.width}; MODEL: {args.model}; PROMPT: {args.prompt}'
+    print(f'GENERATING: {path}; {metadata}')
     generator.manual_seed(seed)
     image = pipe(args.prompt, generator=generator, guidance_scale=args.guidance, num_inference_steps=args.steps, height=args.height, width=args.width).images[0]
-    path = save_path(n)
     image.save(path)
-    metadata = f'seed: {seed}; guidance: {args.guidance}; height: {args.height}; width: {args.width}; model: {args.model}; prompt: {args.prompt}'
-    print(f'saved {path} – metadata: {metadata}')
     with ExifToolHelper() as et:
         et.set_tags([path], tags={'Comment': metadata}, params=['-overwrite_original'])