From 7947b76518f4bf9f0c4db21d6d0b28832c2188c9 Mon Sep 17 00:00:00 2001 From: Christian Heller Date: Fri, 16 Aug 2024 09:12:43 +0200 Subject: [PATCH] Improve readability of output, reduce diffusers logging clutter. --- test.py | 29 +++++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/test.py b/test.py index ef64c44..d485214 100755 --- 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']) -- 2.30.2