home · contact · privacy
Rename executable.
authorChristian Heller <c.heller@plomlompom.de>
Fri, 16 Aug 2024 07:15:34 +0000 (09:15 +0200)
committerChristian Heller <c.heller@plomlompom.de>
Fri, 16 Aug 2024 07:15:34 +0000 (09:15 +0200)
stable.py [new file with mode: 0755]
test.py [deleted file]

diff --git a/stable.py b/stable.py
new file mode 100755 (executable)
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 (executable)
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'])