DiffusionAi/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2025-01-09 04:18:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'development' into webui-integration

Browse Source
This commit is contained in:
Lincoln Stein 2022-09-28 14:21:00 -04:00 committed by GitHub
commit 6e54f504e7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 502 additions and 180 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

79
docs/features/CLI.md
View File

@ -205,6 +205,85 @@ well as the --mask (-M) argument:
| --init_mask <path> | -M<path> | None |Path to an image the same size as the initial_image, with areas for inpainting made transparent.|
# Convenience commands
In addition to the standard image generation arguments, there are a
series of convenience commands that begin with !:
## !fix
This command runs a post-processor on a previously-generated image. It
takes a PNG filename or path and applies your choice of the -U, -G, or
--embiggen switches in order to fix faces or upscale. If you provide a
filename, the script will look for it in the current output
directory. Otherwise you can provide a full or partial path to the
desired file.
Some examples:
Upscale to 4X its original size and fix faces using codeformer:
~~~
dream> !fix 0000045.4829112.png -G1 -U4 -ft codeformer
~~~
Use the GFPGAN algorithm to fix faces, then upscale to 3X using --embiggen:
~~~
dream> !fix 0000045.4829112.png -G0.8 -ft gfpgan
>> fixing outputs/img-samples/0000045.4829112.png
>> retrieved seed 4829112 and prompt "boy enjoying a banana split"
>> GFPGAN - Restoring Faces for image seed:4829112
Outputs:
[1] outputs/img-samples/000017.4829112.gfpgan-00.png: !fix "outputs/img-samples/0000045.4829112.png" -s 50 -S -W 512 -H 512 -C 7.5 -A k_lms -G 0.8
dream> !fix 000017.4829112.gfpgan-00.png --embiggen 3
...lots of text...
Outputs:
[2] outputs/img-samples/000018.2273800735.embiggen-00.png: !fix "outputs/img-samples/000017.243781548.gfpgan-00.png" -s 50 -S 2273800735 -W 512 -H 512 -C 7.5 -A k_lms --embiggen 3.0 0.75 0.25
~~~
## !fetch
This command retrieves the generation parameters from a previously
generated image and either loads them into the command line
(Linux|Mac), or prints them out in a comment for copy-and-paste
(Windows). You may provide either the name of a file in the current
output directory, or a full file path.
~~~
dream> !fetch 0000015.8929913.png
# the script returns the next line, ready for editing and running:
dream> a fantastic alien landscape -W 576 -H 512 -s 60 -A plms -C 7.5
~~~
Note that this command may behave unexpectedly if given a PNG file that
was not generated by InvokeAI.
## !history
The dream script keeps track of all the commands you issue during a
session, allowing you to re-run them. On Mac and Linux systems, it
also writes the command-line history out to disk, giving you access to
the most recent 1000 commands issued.
The `!history` command will return a numbered list of all the commands
issued during the session (Windows), or the most recent 1000 commands
(Mac|Linux). You can then repeat a command by using the command !NNN,
where "NNN" is the history line number. For example:
~~~
dream> !history
...
[14] happy woman sitting under tree wearing broad hat and flowing garment
[15] beautiful woman sitting under tree wearing broad hat and flowing garment
[18] beautiful woman sitting under tree wearing broad hat and flowing garment -v0.2 -n6
[20] watercolor of beautiful woman sitting under tree wearing broad hat and flowing garment -v0.2 -n6 -S2878767194
[21] surrealist painting of beautiful woman sitting under tree wearing broad hat and flowing garment -v0.2 -n6 -S2878767194
...
dream> !20
dream> watercolor of beautiful woman sitting under tree wearing broad hat and flowing garment -v0.2 -n6 -S2878767194
~~~
# Command-line editing and completion
If you are on a Macintosh or Linux machine, the command-line offers

43
ldm/dream/args.py
View File

@ -81,13 +81,14 @@ with metadata_from_png():
"""
import argparse
from argparse import Namespace
from argparse import Namespace, RawTextHelpFormatter
import shlex
import json
import hashlib
import os
import copy
import base64
import functools
import ldm.dream.pngwriter
from ldm.dream.conditioning import split_weighted_subprompts
@ -220,9 +221,15 @@ class Args(object):
# outpainting parameters
if a['out_direction']:
switches.append(f'-D {" ".join([str(u) for u in a["out_direction"]])}')
# LS: slight semantic drift which needs addressing in the future:
# 1. Variations come out of the stored metadata as a packed string with the keyword "variations"
# 2. However, they come out of the CLI (and probably web) with the keyword "with_variations" and
# in broken-out form. Variation (1) should be changed to comply with (2)
if a['with_variations']:
formatted_variations = ','.join(f'{seed}:{weight}' for seed, weight in (a["with_variations"]))
switches.append(f'-V {formatted_variations}')
formatted_variations = ','.join(f'{seed}:{weight}' for seed, weight in (a["variations"]))
switches.append(f'-V {a["formatted_variations"]}')
if 'variations' in a:
switches.append(f'-V {a["variations"]}')
return ' '.join(switches)
def __getattribute__(self,name):
@ -455,9 +462,24 @@ class Args(object):
# This creates the parser that processes commands on the dream> command line
def _create_dream_cmd_parser(self):
parser = argparse.ArgumentParser(
description="""
Generate example: dream> a fantastic alien landscape -W576 -H512 -s60 -n4
Postprocess example: dream> !pp 0000045.4829112.png -G1 -U4 -ft codeformer
formatter_class=RawTextHelpFormatter,
description=
"""
*Image generation:*
dream> a fantastic alien landscape -W576 -H512 -s60 -n4
*postprocessing*
!fix applies upscaling/facefixing to a previously-generated image.
dream> !fix 0000045.4829112.png -G1 -U4 -ft codeformer
*History manipulation*
!fetch retrieves the command used to generate an earlier image.
dream> !fetch 0000015.8929913.png
dream> a fantastic alien landscape -W 576 -H 512 -s 60 -A plms -C 7.5
!history lists all the commands issued during the current session.
!NN retrieves the NNth command from the history
"""
)
render_group = parser.add_argument_group('General rendering')
@ -625,7 +647,7 @@ class Args(object):
'-embiggen',
nargs='+',
type=float,
help='Embiggen tiled img2img for higher resolution and detail without extra VRAM usage. Takes scale factor relative to the size of the --init_img (-I), followed by ESRGAN upscaling strength (0-1.0), followed by minimum amount of overlap between tiles as a decimal ratio (0 - 1.0) or number of pixels. ESRGAN strength defaults to 0.75, and overlap defaults to 0.25 . ESRGAN is used to upscale the init prior to cutting it into tiles/pieces to run through img2img and then stitch back togeather.',
help='Arbitrary upscaling using img2img. Provide scale factor (0.75), optionally followed by strength (0.75) and tile overlap proportion (0.25).',
default=None,
)
postprocessing_group.add_argument(
@ -633,7 +655,7 @@ class Args(object):
'-embiggen_tiles',
nargs='+',
type=int,
help='If while doing Embiggen we are altering only parts of the image, takes a list of tiles by number to process and replace onto the image e.g. `1 3 5`, useful for redoing problematic spots from a prior Embiggen run',
help='For embiggen, provide list of tiles to process and replace onto the image e.g. `1 3 5`.',
default=None,
)
special_effects_group.add_argument(
@ -749,6 +771,7 @@ def metadata_dumps(opt,
return metadata
@functools.lru_cache(maxsize=50)
def metadata_from_png(png_file_path):
'''
Given the path to a PNG file created by dream.py, retrieves
@ -758,6 +781,10 @@ def metadata_from_png(png_file_path):
opts = metadata_loads(meta)
return opts[0]
def dream_cmd_from_png(png_file_path):
opt = metadata_from_png(png_file_path)
return opt.dream_prompt_str()
def metadata_loads(metadata):
'''
Takes the dictionary corresponding to RFC266 (https://github.com/lstein/stable-diffusion/issues/266)

15
ldm/dream/log.py
View File

@ -22,12 +22,17 @@ def write_log(results, log_path, file_types, output_cntr):
def write_log_message(results, output_cntr):
"""logs to the terminal"""
if len(results) == 0:
return output_cntr
log_lines = [f"{path}: {prompt}\n" for path, prompt in results]
for l in log_lines:
output_cntr += 1
print(f"[{output_cntr}] {l}", end="")
return output_cntr
if len(log_lines)>1:
subcntr = 1
for l in log_lines:
print(f"[{output_cntr}.{subcntr}] {l}", end="")
subcntr += 1
else:
print(f"[{output_cntr}] {log_lines[0]}", end="")
return output_cntr+1
def write_log_files(results, log_path, file_types):
for file_type in file_types:

288
ldm/dream/readline.py
View File

@ -1,38 +1,96 @@
"""
Readline helper functions for dream.py (linux and mac only).
You may import the global singleton `completer` to get access to the
completer object itself. This is useful when you want to autocomplete
seeds:
from ldm.dream.readline import completer
completer.add_seed(18247566)
completer.add_seed(9281839)
"""
import os
import re
import atexit
completer = None
# ---------------readline utilities---------------------
try:
import readline
readline_available = True
except:
readline_available = False
#to simulate what happens on windows systems, uncomment
# this line
#readline_available = False
IMG_EXTENSIONS = ('.png','.jpg','.jpeg')
COMMANDS = (
'--steps','-s',
'--seed','-S',
'--iterations','-n',
'--width','-W','--height','-H',
'--cfg_scale','-C',
'--grid','-g',
'--individual','-i',
'--init_img','-I',
'--init_mask','-M',
'--init_color',
'--strength','-f',
'--variants','-v',
'--outdir','-o',
'--sampler','-A','-m',
'--embedding_path',
'--device',
'--grid','-g',
'--gfpgan_strength','-G',
'--upscale','-U',
'-save_orig','--save_original',
'--skip_normalize','-x',
'--log_tokenization','-t',
'!fix','!fetch',
)
IMG_PATH_COMMANDS = (
'--init_img[=\s]','-I',
'--init_mask[=\s]','-M',
'--init_color[=\s]',
'--embedding_path[=\s]',
'--outdir[=\s]'
)
IMG_FILE_COMMANDS=(
'!fix',
'!fetch',
)
path_regexp = '('+'|'.join(IMG_PATH_COMMANDS+IMG_FILE_COMMANDS) + ')\s*\S*$'
class Completer:
def __init__(self, options):
self.options = sorted(options)
self.options = sorted(options)
self.seeds = set()
self.matches = list()
self.default_dir = None
self.linebuffer = None
return
def complete(self, text, state):
'''
Completes dream command line.
BUG: it doesn't correctly complete files that have spaces in the name.
'''
buffer = readline.get_line_buffer()
if text.startswith(('-I', '--init_img','-M','--init_mask',
'--init_color')):
return self._path_completions(text, state, ('.png','.jpg','.jpeg'))
if buffer.strip().endswith('pp') or text.startswith(('.', '/')):
return self._path_completions(text, state, ('.png','.jpg','.jpeg'))
response = None
if state == 0:
if re.search(path_regexp,buffer):
do_shortcut = re.search('^'+'|'.join(IMG_FILE_COMMANDS),buffer)
self.matches = self._path_completions(text, state, IMG_EXTENSIONS,shortcut_ok=do_shortcut)
# looking for a seed
elif re.search('(-S\s*|--seed[=\s])\d*$',buffer):
self.matches= self._seed_completions(text,state)
# This is the first time for this text, so build a match list.
if text:
elif text:
self.matches = [
s for s in self.options if s and s.startswith(text)
]
@ -47,81 +105,164 @@ class Completer:
response = None
return response
def _path_completions(self, text, state, extensions):
# get the path so far
# TODO: replace this mess with a regular expression match
if text.startswith('-I'):
path = text.replace('-I', '', 1).lstrip()
elif text.startswith('--init_img='):
path = text.replace('--init_img=', '', 1).lstrip()
elif text.startswith('--init_mask='):
path = text.replace('--init_mask=', '', 1).lstrip()
elif text.startswith('-M'):
path = text.replace('-M', '', 1).lstrip()
elif text.startswith('--init_color='):
path = text.replace('--init_color=', '', 1).lstrip()
def add_history(self,line):
'''
Pass thru to readline
'''
readline.add_history(line)
def remove_history_item(self,pos):
readline.remove_history_item(pos)
def add_seed(self, seed):
'''
Add a seed to the autocomplete list for display when -S is autocompleted.
'''
if seed is not None:
self.seeds.add(str(seed))
def set_default_dir(self, path):
self.default_dir=path
def get_line(self,index):
try:
line = self.get_history_item(index)
except IndexError:
return None
return line
def get_current_history_length(self):
return readline.get_current_history_length()
def get_history_item(self,index):
return readline.get_history_item(index)
def show_history(self):
'''
Print the session history using the pydoc pager
'''
import pydoc
lines = list()
h_len = self.get_current_history_length()
if h_len < 1:
print('<empty history>')
return
for i in range(0,h_len):
lines.append(f'[{i+1}] {self.get_history_item(i+1)}')
pydoc.pager('\n'.join(lines))
def set_line(self,line)->None:
self.linebuffer = line
readline.redisplay()
def _seed_completions(self, text, state):
m = re.search('(-S\s?|--seed[=\s]?)(\d*)',text)
if m:
switch = m.groups()[0]
partial = m.groups()[1]
else:
path = text
switch = ''
partial = text
matches = list()
for s in self.seeds:
if s.startswith(partial):
matches.append(switch+s)
matches.sort()
return matches
path = os.path.expanduser(path)
if len(path) == 0:
matches.append(text + './')
def _pre_input_hook(self):
if self.linebuffer:
readline.insert_text(self.linebuffer)
readline.redisplay()
self.linebuffer = None
def _path_completions(self, text, state, extensions, shortcut_ok=False):
# separate the switch from the partial path
match = re.search('^(-\w|--\w+=?)(.*)',text)
if match is None:
switch = None
partial_path = text
else:
switch,partial_path = match.groups()
partial_path = partial_path.lstrip()
matches = list()
path = os.path.expanduser(partial_path)
if os.path.isdir(path):
dir = path
elif os.path.dirname(path) != '':
dir = os.path.dirname(path)
dir_list = os.listdir(dir)
for n in dir_list:
if n.startswith('.') and len(n) > 1:
continue
full_path = os.path.join(dir, n)
if full_path.startswith(path):
if os.path.isdir(full_path):
matches.append(
os.path.join(os.path.dirname(text), n) + '/'
)
elif n.endswith(extensions):
matches.append(os.path.join(os.path.dirname(text), n))
else:
dir = ''
path= os.path.join(dir,path)
try:
response = matches[state]
except IndexError:
response = None
return response
dir_list = os.listdir(dir or '.')
if shortcut_ok and os.path.exists(self.default_dir) and dir=='':
dir_list += os.listdir(self.default_dir)
for node in dir_list:
if node.startswith('.') and len(node) > 1:
continue
full_path = os.path.join(dir, node)
if not (node.endswith(extensions) or os.path.isdir(full_path)):
continue
if not full_path.startswith(path):
continue
if switch is None:
match_path = os.path.join(dir,node)
matches.append(match_path+'/' if os.path.isdir(full_path) else match_path)
elif os.path.isdir(full_path):
matches.append(
switch+os.path.join(os.path.dirname(full_path), node) + '/'
)
elif node.endswith(extensions):
matches.append(
switch+os.path.join(os.path.dirname(full_path), node)
)
return matches
class DummyCompleter(Completer):
def __init__(self,options):
super().__init__(options)
self.history = list()
def add_history(self,line):
self.history.append(line)
def get_current_history_length(self):
return len(self.history)
def get_history_item(self,index):
return self.history[index-1]
def remove_history_item(self,index):
return self.history.pop(index-1)
def set_line(self,line):
print(f'# {line}')
if readline_available:
completer = Completer(COMMANDS)
readline.set_completer(
Completer(
[
'--steps','-s',
'--seed','-S',
'--iterations','-n',
'--width','-W','--height','-H',
'--cfg_scale','-C',
'--grid','-g',
'--individual','-i',
'--init_img','-I',
'--init_mask','-M',
'--init_color',
'--strength','-f',
'--variants','-v',
'--outdir','-o',
'--sampler','-A','-m',
'--embedding_path',
'--device',
'--grid','-g',
'--gfpgan_strength','-G',
'--upscale','-U',
'-save_orig','--save_original',
'--skip_normalize','-x',
'--log_tokenization','t',
]
).complete
completer.complete
)
readline.set_auto_history(False)
readline.set_pre_input_hook(completer._pre_input_hook)
readline.set_completer_delims(' ')
readline.parse_and_bind('tab: complete')
readline.parse_and_bind('set print-completions-horizontally off')
readline.parse_and_bind('set page-completions on')
readline.parse_and_bind('set skip-completed-text on')
readline.parse_and_bind('set bell-style visible')
readline.parse_and_bind('set show-all-if-ambiguous on')
histfile = os.path.join(os.path.expanduser('~'), '.dream_history')
try:
readline.read_history_file(histfile)
@ -129,3 +270,6 @@ if readline_available:
except FileNotFoundError:
pass
atexit.register(readline.write_history_file, histfile)
else:
completer = DummyCompleter(COMMANDS)

32
ldm/generate.py
View File

@ -490,25 +490,26 @@ class Generate:
opt = None,
):
# retrieve the seed from the image;
# note that we will try both the new way and the old way, since not all files have the
# metadata (yet)
seed = None
image_metadata = None
prompt = None
try:
args = metadata_from_png(image_path)
seed = args.seed
prompt = args.prompt
print(f'>> retrieved seed {seed} and prompt "{prompt}" from {image_path}')
except:
m = re.search('(\d+)\.png$',image_path)
if m:
seed = m.group(1)
args = metadata_from_png(image_path)
seed = args.seed
prompt = args.prompt
print(f'>> retrieved seed {seed} and prompt "{prompt}" from {image_path}')
if not seed:
print('* Could not recover seed for image. Replacing with 42. This will not affect image quality')
seed = 42
# try to reuse the same filename prefix as the original file.
# note that this is hacky
prefix = None
m = re.search('(\d+)\.',os.path.basename(image_path))
if m:
prefix = m.groups()[0]
# face fixers and esrgan take an Image, but embiggen takes a path
image = Image.open(image_path)
@ -530,6 +531,7 @@ class Generate:
save_original = save_original,
upscale = upscale,
image_callback = callback,
prefix = prefix,
)
elif tool == 'embiggen':
@ -716,7 +718,9 @@ class Generate:
strength = 0.0,
codeformer_fidelity = 0.75,
save_original = False,
image_callback = None):
image_callback = None,
prefix = None,
):
for r in image_list:
image, seed = r
@ -750,7 +754,7 @@ class Generate:
)
if image_callback is not None:
image_callback(image, seed, upscaled=True)
image_callback(image, seed, upscaled=True, use_prefix=prefix)
else:
r[0] = image

225
scripts/dream.py Executable file → Normal file
View File

@ -9,19 +9,16 @@ import copy
import warnings
import time
sys.path.append('.') # corrects a weird problem on Macs
import ldm.dream.readline
from ldm.dream.args import Args, metadata_dumps, metadata_from_png
from ldm.dream.readline import completer
from ldm.dream.args import Args, metadata_dumps, metadata_from_png, dream_cmd_from_png
from ldm.dream.pngwriter import PngWriter
from ldm.dream.image_util import make_grid
from ldm.dream.log import write_log
from omegaconf import OmegaConf
from backend.invoke_ai_web_server import InvokeAIWebServer
# Placeholder to be replaced with proper class that tracks the
# outputs and associates with the prompt that generated them.
# Just want to get the formatting look right for now.
output_cntr = 0
# The output counter labels each output and is keyed to the
# command-line history
output_cntr = completer.get_current_history_length()+1
def main():
"""Initialize command-line parsers and the diffusion model"""
@ -142,7 +139,10 @@ def main_loop(gen, opt, infile):
while not done:
operation = 'generate' # default operation, alternative is 'postprocess'
if completer:
completer.set_default_dir(opt.outdir)
try:
command = get_next_command(infile)
except EOFError:
@ -160,16 +160,28 @@ def main_loop(gen, opt, infile):
done = True
break
if command.startswith(
'!dream'
): # in case a stored prompt still contains the !dream command
if command.startswith('!dream'): # in case a stored prompt still contains the !dream command
command = command.replace('!dream ','',1)
if command.startswith(
'!fix'
):
if command.startswith('!fix'):
command = command.replace('!fix ','',1)
operation = 'postprocess'
if command.startswith('!fetch'):
file_path = command.replace('!fetch ','',1)
retrieve_dream_command(opt,file_path)
continue
if command == '!history':
completer.show_history()
continue
match = re.match('^!(\d+)',command)
if match:
command_no = match.groups()[0]
command = completer.get_line(int(command_no))
completer.set_line(command)
continue
if opt.parse_cmd(command) is None:
continue
@ -220,37 +232,15 @@ def main_loop(gen, opt, infile):
opt.strength = 0.75 if opt.out_direction is None else 0.83
if opt.with_variations is not None:
# shotgun parsing, woo
parts = []
broken = False # python doesn't have labeled loops...
for part in opt.with_variations.split(','):
seed_and_weight = part.split(':')
if len(seed_and_weight) != 2:
print(f'could not parse with_variation part "{part}"')
broken = True
break
try:
seed = int(seed_and_weight[0])
weight = float(seed_and_weight[1])
except ValueError:
print(f'could not parse with_variation part "{part}"')
broken = True
break
parts.append([seed, weight])
if broken:
continue
if len(parts) > 0:
opt.with_variations = parts
else:
opt.with_variations = None
opt.with_variations = split_variations(opt.with_variations)
if opt.prompt_as_dir:
# sanitize the prompt to a valid folder name
subdir = path_filter.sub('_', opt.prompt)[:name_max].rstrip(' .')
# truncate path to maximum allowed length
# 27 is the length of '######.##########.##.png', plus two separators and a NUL
subdir = subdir[:(path_max - 27 - len(os.path.abspath(opt.outdir)))]
# 39 is the length of '######.##########.##########-##.png', plus two separators and a NUL
subdir = subdir[:(path_max - 39 - len(os.path.abspath(opt.outdir)))]
current_outdir = os.path.join(opt.outdir, subdir)
print('Writing files to directory: "' + current_outdir + '"')
@ -267,37 +257,35 @@ def main_loop(gen, opt, infile):
last_results = []
try:
file_writer = PngWriter(current_outdir)
prefix = file_writer.unique_prefix()
results = [] # list of filename, prompt pairs
grid_images = dict() # seed -> Image, only used if `opt.grid`
prior_variations = opt.with_variations or []
def image_writer(image, seed, upscaled=False, first_seed=None):
def image_writer(image, seed, upscaled=False, first_seed=None, use_prefix=None):
# note the seed is the seed of the current image
# the first_seed is the original seed that noise is added to
# when the -v switch is used to generate variations
path = None
nonlocal prior_variations
if use_prefix is not None:
prefix = use_prefix
else:
prefix = file_writer.unique_prefix()
path = None
if opt.grid:
grid_images[seed] = image
else:
if operation == 'postprocess':
filename = choose_postprocess_name(opt.prompt)
elif upscaled and opt.save_original:
filename = f'{prefix}.{seed}.postprocessed.png'
else:
filename = f'{prefix}.{seed}.png'
if opt.variation_amount > 0:
first_seed = first_seed or seed
this_variation = [[seed, opt.variation_amount]]
opt.with_variations = prior_variations + this_variation
formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed)
elif len(prior_variations) > 0:
formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed)
elif operation == 'postprocess':
formatted_dream_prompt = '!fix '+opt.dream_prompt_str(seed=seed)
else:
formatted_dream_prompt = opt.dream_prompt_str(seed=seed)
postprocessed = upscaled if upscaled else operation=='postprocess'
filename, formatted_dream_prompt = prepare_image_metadata(
opt,
prefix,
seed,
operation,
prior_variations,
postprocessed,
first_seed
)
path = file_writer.save_image_and_prompt_to_png(
image = image,
dream_prompt = formatted_dream_prompt,
@ -311,10 +299,15 @@ def main_loop(gen, opt, infile):
if (not upscaled) or opt.save_original:
# only append to results if we didn't overwrite an earlier output
results.append([path, formatted_dream_prompt])
# so that the seed autocompletes (on linux|mac when -S or --seed specified
if completer:
completer.add_seed(seed)
completer.add_seed(first_seed)
last_results.append([path, seed])
if operation == 'generate':
catch_ctrl_c = infile is None # if running interactively, we catch keyboard interrupts
opt.last_operation='generate'
gen.prompt2image(
image_callback=image_writer,
catch_interrupts=catch_ctrl_c,
@ -322,7 +315,7 @@ def main_loop(gen, opt, infile):
)
elif operation == 'postprocess':
print(f'>> fixing {opt.prompt}')
do_postprocess(gen,opt,image_writer)
opt.last_operation = do_postprocess(gen,opt,image_writer)
if opt.grid and len(grid_images) > 0:
grid_img = make_grid(list(grid_images.values()))
@ -357,6 +350,10 @@ def main_loop(gen, opt, infile):
global output_cntr
output_cntr = write_log(results, log_path ,('txt', 'md'), output_cntr)
print()
if operation == 'postprocess':
completer.add_history(f'!fix {command}')
else:
completer.add_history(command)
print('goodbye!')
@ -378,8 +375,9 @@ def do_postprocess (gen, opt, callback):
elif opt.out_direction:
tool = 'outpaint'
opt.save_original = True # do not overwrite old image!
return gen.apply_postprocessor(
image_path = opt.prompt,
opt.last_operation = f'postprocess:{tool}'
gen.apply_postprocessor(
image_path = file_path,
tool = tool,
gfpgan_strength = opt.gfpgan_strength,
codeformer_fidelity = opt.codeformer_fidelity,
@ -389,18 +387,54 @@ def do_postprocess (gen, opt, callback):
callback = callback,
opt = opt,
)
return opt.last_operation
def choose_postprocess_name(original_filename):
basename,_ = os.path.splitext(os.path.basename(original_filename))
if re.search('\d+\.\d+$',basename):
return f'{basename}.fixed.png'
match = re.search('(\d+\.\d+)\.fixed(-(\d+))?$',basename)
if match:
counter = match.group(3) or 0
return '{prefix}-{counter:02d}.png'.format(prefix=match.group(1), counter=int(counter)+1)
def prepare_image_metadata(
opt,
prefix,
seed,
operation='generate',
prior_variations=[],
postprocessed=False,
first_seed=None
):
if postprocessed and opt.save_original:
filename = choose_postprocess_name(opt,prefix,seed)
else:
return f'{basename}.fixed.png'
filename = f'{prefix}.{seed}.png'
if opt.variation_amount > 0:
first_seed = first_seed or seed
this_variation = [[seed, opt.variation_amount]]
opt.with_variations = prior_variations + this_variation
formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed)
elif len(prior_variations) > 0:
formatted_dream_prompt = opt.dream_prompt_str(seed=first_seed)
elif operation == 'postprocess':
formatted_dream_prompt = '!fix '+opt.dream_prompt_str(seed=seed)
else:
formatted_dream_prompt = opt.dream_prompt_str(seed=seed)
return filename,formatted_dream_prompt
def choose_postprocess_name(opt,prefix,seed) -> str:
match = re.search('postprocess:(\w+)',opt.last_operation)
if match:
modifier = match.group(1) # will look like "gfpgan", "upscale", "outpaint" or "embiggen"
else:
modifier = 'postprocessed'
counter = 0
filename = None
available = False
while not available:
if counter == 0:
filename = f'{prefix}.{seed}.{modifier}.png'
else:
filename = f'{prefix}.{seed}.{modifier}-{counter:02d}.png'
available = not os.path.exists(os.path.join(opt.outdir,filename))
counter += 1
return filename
def get_next_command(infile=None) -> str: # command string
if infile is None:
@ -430,16 +464,45 @@ def invoke_ai_web_server_loop(gen, gfpgan, codeformer, esrgan):
pass
def write_log_message(results, log_path):
"""logs the name of the output image, prompt, and prompt args to the terminal and log file"""
global output_cntr
log_lines = [f'{path}: {prompt}\n' for path, prompt in results]
for l in log_lines:
output_cntr += 1
print(f'[{output_cntr}] {l}',end='')
def split_variations(variations_string) -> list:
# shotgun parsing, woo
parts = []
broken = False # python doesn't have labeled loops...
for part in variations_string.split(','):
seed_and_weight = part.split(':')
if len(seed_and_weight) != 2:
print(f'** Could not parse with_variation part "{part}"')
broken = True
break
try:
seed = int(seed_and_weight[0])
weight = float(seed_and_weight[1])
except ValueError:
print(f'** Could not parse with_variation part "{part}"')
broken = True
break
parts.append([seed, weight])
if broken:
return None
elif len(parts) == 0:
return None
else:
return parts
with open(log_path, 'a', encoding='utf-8') as file:
file.writelines(log_lines)
def retrieve_dream_command(opt,file_path):
'''
Given a full or partial path to a previously-generated image file,
will retrieve and format the dream command used to generate the image,
and pop it into the readline buffer (linux, Mac), or print out a comment
for cut-and-paste (windows)
'''
dir,basename = os.path.split(file_path)
if len(dir) == 0:
path = os.path.join(opt.outdir,basename)
else:
path = file_path
cmd = dream_cmd_from_png(path)
completer.set_line(cmd)
if __name__ == '__main__':
main()