From: Christian Heller Date: Fri, 16 Aug 2024 07:15:34 +0000 (+0200) Subject: Rename executable. X-Git-Url: https://plomlompom.com/repos/?a=commitdiff_plain;h=2b9f67c2835a23ee4f90aaa8a30036599dbdab72;p=stable_plom Rename executable. --- diff --git a/stable.py b/stable.py new file mode 100755 index 0000000..d485214 --- /dev/null +++ b/stable.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +from argparse import ArgumentParser +from random import randint +from PIL import Image +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) + parser.add_argument('-o', '--output', required=True) + parser.add_argument('-r', '--randomness_seed', default=1, type=int, help='default 1; if set 0, chosen randomnly') + parser.add_argument('-g', '--guidance', default=7.5, type=float, help='default 7.5') + parser.add_argument('-s', '--steps', default=15, type=int, help='default 15') + parser.add_argument('-m', '--model', default=DEFAULT_MODEL, type=str, help=f'default: {DEFAULT_MODEL}') + parser.add_argument('-h', '--height', default=512, type=int, help='default: 512') + parser.add_argument('-w', '--width', default=512, type=int, help='default=512') + parser.add_argument('-n', '--number', default=1, type=int, help='default=1') + parser.add_argument('-H', '--help', action='help') + return parser.parse_args() + +args = parse_args() + +dir_path = dirname(args.output) +filename = basename(args.output) +filename_sans_ext, ext = splitext(filename) +ext = ext[1:] if ext else 'png' + +def save_path(n: int) -> str: + filename_count = f'_{n:08}' if args.number > 1 else '' + return path_join(dir_path, f'{filename_sans_ext}{filename_count}.{ext}') + +for n in range(args.number): + path = save_path(n) + if exists(path): + raise Exception(f'Would overwrite file: {path}') + +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] + image.save(path) + with ExifToolHelper() as et: + et.set_tags([path], tags={'Comment': metadata}, params=['-overwrite_original']) diff --git a/test.py b/test.py deleted file mode 100755 index d485214..0000000 --- a/test.py +++ /dev/null @@ -1,76 +0,0 @@ -#!/usr/bin/env python3 -from argparse import ArgumentParser -from random import randint -from PIL import Image -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) - parser.add_argument('-o', '--output', required=True) - parser.add_argument('-r', '--randomness_seed', default=1, type=int, help='default 1; if set 0, chosen randomnly') - parser.add_argument('-g', '--guidance', default=7.5, type=float, help='default 7.5') - parser.add_argument('-s', '--steps', default=15, type=int, help='default 15') - parser.add_argument('-m', '--model', default=DEFAULT_MODEL, type=str, help=f'default: {DEFAULT_MODEL}') - parser.add_argument('-h', '--height', default=512, type=int, help='default: 512') - parser.add_argument('-w', '--width', default=512, type=int, help='default=512') - parser.add_argument('-n', '--number', default=1, type=int, help='default=1') - parser.add_argument('-H', '--help', action='help') - return parser.parse_args() - -args = parse_args() - -dir_path = dirname(args.output) -filename = basename(args.output) -filename_sans_ext, ext = splitext(filename) -ext = ext[1:] if ext else 'png' - -def save_path(n: int) -> str: - filename_count = f'_{n:08}' if args.number > 1 else '' - return path_join(dir_path, f'{filename_sans_ext}{filename_count}.{ext}') - -for n in range(args.number): - path = save_path(n) - if exists(path): - raise Exception(f'Would overwrite file: {path}') - -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] - image.save(path) - with ExifToolHelper() as et: - et.set_tags([path], tags={'Comment': metadata}, params=['-overwrite_original'])