Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00: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
parent 984575b579 f93963cd6b
commit 6e54f504e7
6 changed files with 502 additions and 180 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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()