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)
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'])