mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Add back old dream.py as legacy_api.py
This commit "reverts" the new API changes by extracting the old functionality into new files. The work is based on the commit `803a51d5adca7e6e28491fc414fd3937bee7cb79` PngWriter regained PromptFormatter as old server used that. `server_legacy.py` is the old server that `dream.py` used. Finally `legacy_api.py` is what `dream.py` used to be at the mentioned commit. One manually run test has been added in order to be able to test compatibility with the old API, currently just testing that the API endpoint works the same way + the image hash is the same as it used to be before.
This commit is contained in:
CapableWeb
Lincoln Stein
parent
ca6385e6fa
commit
6c0dd9b5ef
@ -66,3 +66,43 @@ def write_metadata(img_path:str, meta:dict):
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text('sd-metadata', json.dumps(meta))
|
||||
im.save(img_path,'PNG',pnginfo=info)
|
||||
|
||||
class PromptFormatter:
|
||||
def __init__(self, t2i, opt):
|
||||
self.t2i = t2i
|
||||
self.opt = opt
|
||||
|
||||
# note: the t2i object should provide all these values.
|
||||
# there should be no need to or against opt values
|
||||
def normalize_prompt(self):
|
||||
"""Normalize the prompt and switches"""
|
||||
t2i = self.t2i
|
||||
opt = self.opt
|
||||
|
||||
switches = list()
|
||||
switches.append(f'"{opt.prompt}"')
|
||||
switches.append(f'-s{opt.steps or t2i.steps}')
|
||||
switches.append(f'-W{opt.width or t2i.width}')
|
||||
switches.append(f'-H{opt.height or t2i.height}')
|
||||
switches.append(f'-C{opt.cfg_scale or t2i.cfg_scale}')
|
||||
switches.append(f'-A{opt.sampler_name or t2i.sampler_name}')
|
||||
# to do: put model name into the t2i object
|
||||
# switches.append(f'--model{t2i.model_name}')
|
||||
if opt.seamless or t2i.seamless:
|
||||
switches.append(f'--seamless')
|
||||
if opt.init_img:
|
||||
switches.append(f'-I{opt.init_img}')
|
||||
if opt.fit:
|
||||
switches.append(f'--fit')
|
||||
if opt.strength and opt.init_img is not None:
|
||||
switches.append(f'-f{opt.strength or t2i.strength}')
|
||||
if opt.gfpgan_strength:
|
||||
switches.append(f'-G{opt.gfpgan_strength}')
|
||||
if opt.upscale:
|
||||
switches.append(f'-U {" ".join([str(u) for u in opt.upscale])}')
|
||||
if opt.variation_amount > 0:
|
||||
switches.append(f'-v{opt.variation_amount}')
|
||||
if opt.with_variations:
|
||||
formatted_variations = ','.join(f'{seed}:{weight}' for seed, weight in opt.with_variations)
|
||||
switches.append(f'-V{formatted_variations}')
|
||||
return ' '.join(switches)
|
||||
|
||||
246
ldm/invoke/server_legacy.py
Normal file
246
ldm/invoke/server_legacy.py
Normal file
@ -0,0 +1,246 @@
|
||||
import argparse
|
||||
import json
|
||||
import base64
|
||||
import mimetypes
|
||||
import os
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from ldm.invoke.pngwriter import PngWriter, PromptFormatter
|
||||
from threading import Event
|
||||
|
||||
def build_opt(post_data, seed, gfpgan_model_exists):
|
||||
opt = argparse.Namespace()
|
||||
setattr(opt, 'prompt', post_data['prompt'])
|
||||
setattr(opt, 'init_img', post_data['initimg'])
|
||||
setattr(opt, 'strength', float(post_data['strength']))
|
||||
setattr(opt, 'iterations', int(post_data['iterations']))
|
||||
setattr(opt, 'steps', int(post_data['steps']))
|
||||
setattr(opt, 'width', int(post_data['width']))
|
||||
setattr(opt, 'height', int(post_data['height']))
|
||||
setattr(opt, 'seamless', 'seamless' in post_data)
|
||||
setattr(opt, 'fit', 'fit' in post_data)
|
||||
setattr(opt, 'mask', 'mask' in post_data)
|
||||
setattr(opt, 'invert_mask', 'invert_mask' in post_data)
|
||||
setattr(opt, 'cfg_scale', float(post_data['cfg_scale']))
|
||||
setattr(opt, 'sampler_name', post_data['sampler_name'])
|
||||
setattr(opt, 'gfpgan_strength', float(post_data['gfpgan_strength']) if gfpgan_model_exists else 0)
|
||||
setattr(opt, 'upscale', [int(post_data['upscale_level']), float(post_data['upscale_strength'])] if post_data['upscale_level'] != '' else None)
|
||||
setattr(opt, 'progress_images', 'progress_images' in post_data)
|
||||
setattr(opt, 'seed', None if int(post_data['seed']) == -1 else int(post_data['seed']))
|
||||
setattr(opt, 'variation_amount', float(post_data['variation_amount']) if int(post_data['seed']) != -1 else 0)
|
||||
setattr(opt, 'with_variations', [])
|
||||
|
||||
broken = False
|
||||
if int(post_data['seed']) != -1 and post_data['with_variations'] != '':
|
||||
for part in post_data['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
|
||||
opt.with_variations.append([seed, weight])
|
||||
|
||||
if broken:
|
||||
raise CanceledException
|
||||
|
||||
if len(opt.with_variations) == 0:
|
||||
opt.with_variations = None
|
||||
|
||||
return opt
|
||||
|
||||
class CanceledException(Exception):
|
||||
pass
|
||||
|
||||
class DreamServer(BaseHTTPRequestHandler):
|
||||
model = None
|
||||
outdir = None
|
||||
canceled = Event()
|
||||
|
||||
def do_GET(self):
|
||||
if self.path == "/":
|
||||
self.send_response(200)
|
||||
self.send_header("Content-type", "text/html")
|
||||
self.end_headers()
|
||||
with open("./static/dream_web/index.html", "rb") as content:
|
||||
self.wfile.write(content.read())
|
||||
elif self.path == "/config.js":
|
||||
# unfortunately this import can't be at the top level, since that would cause a circular import
|
||||
from ldm.gfpgan.gfpgan_tools import gfpgan_model_exists
|
||||
self.send_response(200)
|
||||
self.send_header("Content-type", "application/javascript")
|
||||
self.end_headers()
|
||||
config = {
|
||||
'gfpgan_model_exists': gfpgan_model_exists
|
||||
}
|
||||
self.wfile.write(bytes("let config = " + json.dumps(config) + ";\n", "utf-8"))
|
||||
elif self.path == "/run_log.json":
|
||||
self.send_response(200)
|
||||
self.send_header("Content-type", "application/json")
|
||||
self.end_headers()
|
||||
output = []
|
||||
|
||||
log_file = os.path.join(self.outdir, "dream_web_log.txt")
|
||||
if os.path.exists(log_file):
|
||||
with open(log_file, "r") as log:
|
||||
for line in log:
|
||||
url, config = line.split(": {", maxsplit=1)
|
||||
config = json.loads("{" + config)
|
||||
config["url"] = url.lstrip(".")
|
||||
if os.path.exists(url):
|
||||
output.append(config)
|
||||
|
||||
self.wfile.write(bytes(json.dumps({"run_log": output}), "utf-8"))
|
||||
elif self.path == "/cancel":
|
||||
self.canceled.set()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(bytes('{}', 'utf8'))
|
||||
else:
|
||||
path = "." + self.path
|
||||
cwd = os.path.realpath(os.getcwd())
|
||||
is_in_cwd = os.path.commonprefix((os.path.realpath(path), cwd)) == cwd
|
||||
if not (is_in_cwd and os.path.exists(path)):
|
||||
self.send_response(404)
|
||||
return
|
||||
mime_type = mimetypes.guess_type(path)[0]
|
||||
if mime_type is not None:
|
||||
self.send_response(200)
|
||||
self.send_header("Content-type", mime_type)
|
||||
self.end_headers()
|
||||
with open("." + self.path, "rb") as content:
|
||||
self.wfile.write(content.read())
|
||||
else:
|
||||
self.send_response(404)
|
||||
|
||||
def do_POST(self):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-type", "application/json")
|
||||
self.end_headers()
|
||||
|
||||
# unfortunately this import can't be at the top level, since that would cause a circular import
|
||||
# TODO temporarily commented out, import fails for some reason
|
||||
# from ldm.gfpgan.gfpgan_tools import gfpgan_model_exists
|
||||
gfpgan_model_exists = False
|
||||
|
||||
content_length = int(self.headers['Content-Length'])
|
||||
post_data = json.loads(self.rfile.read(content_length))
|
||||
opt = build_opt(post_data, self.model.seed, gfpgan_model_exists)
|
||||
|
||||
self.canceled.clear()
|
||||
print(f">> Request to generate with prompt: {opt.prompt}")
|
||||
# In order to handle upscaled images, the PngWriter needs to maintain state
|
||||
# across images generated by each call to prompt2img(), so we define it in
|
||||
# the outer scope of image_done()
|
||||
config = post_data.copy() # Shallow copy
|
||||
config['initimg'] = config.pop('initimg_name', '')
|
||||
|
||||
images_generated = 0 # helps keep track of when upscaling is started
|
||||
images_upscaled = 0 # helps keep track of when upscaling is completed
|
||||
pngwriter = PngWriter(self.outdir)
|
||||
|
||||
prefix = pngwriter.unique_prefix()
|
||||
# if upscaling is requested, then this will be called twice, once when
|
||||
# the images are first generated, and then again when after upscaling
|
||||
# is complete. The upscaling replaces the original file, so the second
|
||||
# entry should not be inserted into the image list.
|
||||
def image_done(image, seed, upscaled=False, first_seed=-1, use_prefix=None):
|
||||
print(f'First seed: {first_seed}')
|
||||
name = f'{prefix}.{seed}.png'
|
||||
iter_opt = argparse.Namespace(**vars(opt)) # copy
|
||||
if opt.variation_amount > 0:
|
||||
this_variation = [[seed, opt.variation_amount]]
|
||||
if opt.with_variations is None:
|
||||
iter_opt.with_variations = this_variation
|
||||
else:
|
||||
iter_opt.with_variations = opt.with_variations + this_variation
|
||||
iter_opt.variation_amount = 0
|
||||
elif opt.with_variations is None:
|
||||
iter_opt.seed = seed
|
||||
normalized_prompt = PromptFormatter(self.model, iter_opt).normalize_prompt()
|
||||
path = pngwriter.save_image_and_prompt_to_png(image, f'{normalized_prompt} -S{iter_opt.seed}', name)
|
||||
|
||||
if int(config['seed']) == -1:
|
||||
config['seed'] = seed
|
||||
# Append post_data to log, but only once!
|
||||
if not upscaled:
|
||||
with open(os.path.join(self.outdir, "dream_web_log.txt"), "a") as log:
|
||||
log.write(f"{path}: {json.dumps(config)}\n")
|
||||
|
||||
self.wfile.write(bytes(json.dumps(
|
||||
{'event': 'result', 'url': path, 'seed': seed, 'config': config}
|
||||
) + '\n',"utf-8"))
|
||||
|
||||
# control state of the "postprocessing..." message
|
||||
upscaling_requested = opt.upscale or opt.gfpgan_strength > 0
|
||||
nonlocal images_generated # NB: Is this bad python style? It is typical usage in a perl closure.
|
||||
nonlocal images_upscaled # NB: Is this bad python style? It is typical usage in a perl closure.
|
||||
if upscaled:
|
||||
images_upscaled += 1
|
||||
else:
|
||||
images_generated += 1
|
||||
if upscaling_requested:
|
||||
action = None
|
||||
if images_generated >= opt.iterations:
|
||||
if images_upscaled < opt.iterations:
|
||||
action = 'upscaling-started'
|
||||
else:
|
||||
action = 'upscaling-done'
|
||||
if action:
|
||||
x = images_upscaled + 1
|
||||
self.wfile.write(bytes(json.dumps(
|
||||
{'event': action, 'processed_file_cnt': f'{x}/{opt.iterations}'}
|
||||
) + '\n',"utf-8"))
|
||||
|
||||
step_writer = PngWriter(os.path.join(self.outdir, "intermediates"))
|
||||
step_index = 1
|
||||
def image_progress(sample, step):
|
||||
if self.canceled.is_set():
|
||||
self.wfile.write(bytes(json.dumps({'event':'canceled'}) + '\n', 'utf-8'))
|
||||
raise CanceledException
|
||||
path = None
|
||||
# since rendering images is moderately expensive, only render every 5th image
|
||||
# and don't bother with the last one, since it'll render anyway
|
||||
nonlocal step_index
|
||||
if opt.progress_images and step % 5 == 0 and step < opt.steps - 1:
|
||||
image = self.model.sample_to_image(sample)
|
||||
name = f'{prefix}.{opt.seed}.{step_index}.png'
|
||||
metadata = f'{opt.prompt} -S{opt.seed} [intermediate]'
|
||||
path = step_writer.save_image_and_prompt_to_png(image, metadata, name)
|
||||
step_index += 1
|
||||
self.wfile.write(bytes(json.dumps(
|
||||
{'event': 'step', 'step': step + 1, 'url': path}
|
||||
) + '\n',"utf-8"))
|
||||
|
||||
try:
|
||||
if opt.init_img is None:
|
||||
# Run txt2img
|
||||
self.model.prompt2image(**vars(opt), step_callback=image_progress, image_callback=image_done)
|
||||
else:
|
||||
# Decode initimg as base64 to temp file
|
||||
with open("./img2img-tmp.png", "wb") as f:
|
||||
initimg = opt.init_img.split(",")[1] # Ignore mime type
|
||||
f.write(base64.b64decode(initimg))
|
||||
opt1 = argparse.Namespace(**vars(opt))
|
||||
opt1.init_img = "./img2img-tmp.png"
|
||||
|
||||
try:
|
||||
# Run img2img
|
||||
self.model.prompt2image(**vars(opt1), step_callback=image_progress, image_callback=image_done)
|
||||
finally:
|
||||
# Remove the temp file
|
||||
os.remove("./img2img-tmp.png")
|
||||
except CanceledException:
|
||||
print(f"Canceled.")
|
||||
return
|
||||
|
||||
|
||||
class ThreadingDreamServer(ThreadingHTTPServer):
|
||||
def __init__(self, server_address):
|
||||
super(ThreadingDreamServer, self).__init__(server_address, DreamServer)
|
||||
685
scripts/legacy_api.py
Executable file
685
scripts/legacy_api.py
Executable file
@ -0,0 +1,685 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) 2022 Lincoln D. Stein (https://github.com/lstein)
|
||||
|
||||
import argparse
|
||||
import shlex
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import copy
|
||||
import warnings
|
||||
import time
|
||||
import ldm.invoke.readline
|
||||
from ldm.invoke.pngwriter import PngWriter, PromptFormatter
|
||||
from ldm.invoke.server_legacy import DreamServer, ThreadingDreamServer
|
||||
from ldm.invoke.image_util import make_grid
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
# 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
|
||||
|
||||
|
||||
def main():
|
||||
"""Initialize command-line parsers and the diffusion model"""
|
||||
arg_parser = create_argv_parser()
|
||||
opt = arg_parser.parse_args()
|
||||
|
||||
if opt.laion400m:
|
||||
print('--laion400m flag has been deprecated. Please use --model laion400m instead.')
|
||||
sys.exit(-1)
|
||||
if opt.weights != 'model':
|
||||
print('--weights argument has been deprecated. Please configure ./configs/models.yaml, and call it using --model instead.')
|
||||
sys.exit(-1)
|
||||
|
||||
try:
|
||||
models = OmegaConf.load(opt.config)
|
||||
width = models[opt.model].width
|
||||
height = models[opt.model].height
|
||||
config = models[opt.model].config
|
||||
weights = models[opt.model].weights
|
||||
except (FileNotFoundError, IOError, KeyError) as e:
|
||||
print(f'{e}. Aborting.')
|
||||
sys.exit(-1)
|
||||
|
||||
print('* Initializing, be patient...\n')
|
||||
sys.path.append('.')
|
||||
from pytorch_lightning import logging
|
||||
from ldm.generate import Generate
|
||||
|
||||
# these two lines prevent a horrible warning message from appearing
|
||||
# when the frozen CLIP tokenizer is imported
|
||||
import transformers
|
||||
|
||||
transformers.logging.set_verbosity_error()
|
||||
|
||||
# creating a simple text2image object with a handful of
|
||||
# defaults passed on the command line.
|
||||
# additional parameters will be added (or overriden) during
|
||||
# the user input loop
|
||||
t2i = Generate(
|
||||
# width=width,
|
||||
# height=height,
|
||||
sampler_name=opt.sampler_name,
|
||||
weights=weights,
|
||||
full_precision=opt.full_precision,
|
||||
config=config,
|
||||
# grid=opt.grid,
|
||||
# this is solely for recreating the prompt
|
||||
# seamless=opt.seamless,
|
||||
embedding_path=opt.embedding_path,
|
||||
# device_type=opt.device,
|
||||
# ignore_ctrl_c=opt.infile is None,
|
||||
)
|
||||
|
||||
# make sure the output directory exists
|
||||
if not os.path.exists(opt.outdir):
|
||||
os.makedirs(opt.outdir)
|
||||
|
||||
# gets rid of annoying messages about random seed
|
||||
logging.getLogger('pytorch_lightning').setLevel(logging.ERROR)
|
||||
|
||||
# load the infile as a list of lines
|
||||
infile = None
|
||||
if opt.infile:
|
||||
try:
|
||||
if os.path.isfile(opt.infile):
|
||||
infile = open(opt.infile, 'r', encoding='utf-8')
|
||||
elif opt.infile == '-': # stdin
|
||||
infile = sys.stdin
|
||||
else:
|
||||
raise FileNotFoundError(f'{opt.infile} not found.')
|
||||
except (FileNotFoundError, IOError) as e:
|
||||
print(f'{e}. Aborting.')
|
||||
sys.exit(-1)
|
||||
|
||||
if opt.seamless:
|
||||
print(">> changed to seamless tiling mode")
|
||||
|
||||
# preload the model
|
||||
t2i.load_model()
|
||||
|
||||
if not infile:
|
||||
print(
|
||||
"\n* Initialization done! Awaiting your command (-h for help, 'q' to quit)"
|
||||
)
|
||||
|
||||
cmd_parser = create_cmd_parser()
|
||||
if opt.web:
|
||||
dream_server_loop(t2i, opt.host, opt.port, opt.outdir)
|
||||
else:
|
||||
main_loop(t2i, opt.outdir, opt.prompt_as_dir, cmd_parser, infile)
|
||||
|
||||
|
||||
def main_loop(t2i, outdir, prompt_as_dir, parser, infile):
|
||||
"""prompt/read/execute loop"""
|
||||
done = False
|
||||
path_filter = re.compile(r'[<>:"/\\|?*]')
|
||||
last_results = list()
|
||||
|
||||
# os.pathconf is not available on Windows
|
||||
if hasattr(os, 'pathconf'):
|
||||
path_max = os.pathconf(outdir, 'PC_PATH_MAX')
|
||||
name_max = os.pathconf(outdir, 'PC_NAME_MAX')
|
||||
else:
|
||||
path_max = 260
|
||||
name_max = 255
|
||||
|
||||
while not done:
|
||||
try:
|
||||
command = get_next_command(infile)
|
||||
except EOFError:
|
||||
done = True
|
||||
continue
|
||||
except KeyboardInterrupt:
|
||||
done = True
|
||||
continue
|
||||
|
||||
# skip empty lines
|
||||
if not command.strip():
|
||||
continue
|
||||
|
||||
if command.startswith(('#', '//')):
|
||||
continue
|
||||
|
||||
# before splitting, escape single quotes so as not to mess
|
||||
# up the parser
|
||||
command = command.replace("'", "\\'")
|
||||
|
||||
try:
|
||||
elements = shlex.split(command)
|
||||
except ValueError as e:
|
||||
print(str(e))
|
||||
continue
|
||||
|
||||
if elements[0] == 'q':
|
||||
done = True
|
||||
break
|
||||
|
||||
if elements[0].startswith(
|
||||
'!dream'
|
||||
): # in case a stored prompt still contains the !dream command
|
||||
elements.pop(0)
|
||||
|
||||
# rearrange the arguments to mimic how it works in the Dream bot.
|
||||
switches = ['']
|
||||
switches_started = False
|
||||
|
||||
for el in elements:
|
||||
if el[0] == '-' and not switches_started:
|
||||
switches_started = True
|
||||
if switches_started:
|
||||
switches.append(el)
|
||||
else:
|
||||
switches[0] += el
|
||||
switches[0] += ' '
|
||||
switches[0] = switches[0][: len(switches[0]) - 1]
|
||||
|
||||
try:
|
||||
opt = parser.parse_args(switches)
|
||||
except SystemExit:
|
||||
parser.print_help()
|
||||
continue
|
||||
if len(opt.prompt) == 0:
|
||||
print('Try again with a prompt!')
|
||||
continue
|
||||
# retrieve previous value!
|
||||
if opt.init_img is not None and re.match('^-\\d+$', opt.init_img):
|
||||
try:
|
||||
opt.init_img = last_results[int(opt.init_img)][0]
|
||||
print(f'>> Reusing previous image {opt.init_img}')
|
||||
except IndexError:
|
||||
print(
|
||||
f'>> No previous initial image at position {opt.init_img} found')
|
||||
opt.init_img = None
|
||||
continue
|
||||
|
||||
if opt.seed is not None and opt.seed < 0: # retrieve previous value!
|
||||
try:
|
||||
opt.seed = last_results[opt.seed][1]
|
||||
print(f'>> Reusing previous seed {opt.seed}')
|
||||
except IndexError:
|
||||
print(f'>> No previous seed at position {opt.seed} found')
|
||||
opt.seed = None
|
||||
continue
|
||||
|
||||
do_grid = opt.grid or t2i.grid
|
||||
|
||||
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
|
||||
|
||||
if opt.outdir:
|
||||
if not os.path.exists(opt.outdir):
|
||||
os.makedirs(opt.outdir)
|
||||
current_outdir = opt.outdir
|
||||
elif 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(outdir)))]
|
||||
current_outdir = os.path.join(outdir, subdir)
|
||||
|
||||
print('Writing files to directory: "' + current_outdir + '"')
|
||||
|
||||
# make sure the output directory exists
|
||||
if not os.path.exists(current_outdir):
|
||||
os.makedirs(current_outdir)
|
||||
else:
|
||||
current_outdir = outdir
|
||||
|
||||
# Here is where the images are actually generated!
|
||||
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 `do_grid`
|
||||
|
||||
def image_writer(image, seed, upscaled=False):
|
||||
path = None
|
||||
if do_grid:
|
||||
grid_images[seed] = image
|
||||
else:
|
||||
if upscaled and opt.save_original:
|
||||
filename = f'{prefix}.{seed}.postprocessed.png'
|
||||
else:
|
||||
filename = f'{prefix}.{seed}.png'
|
||||
if opt.variation_amount > 0:
|
||||
iter_opt = argparse.Namespace(**vars(opt)) # copy
|
||||
this_variation = [[seed, opt.variation_amount]]
|
||||
if opt.with_variations is None:
|
||||
iter_opt.with_variations = this_variation
|
||||
else:
|
||||
iter_opt.with_variations = opt.with_variations + this_variation
|
||||
iter_opt.variation_amount = 0
|
||||
normalized_prompt = PromptFormatter(
|
||||
t2i, iter_opt).normalize_prompt()
|
||||
metadata_prompt = f'{normalized_prompt} -S{iter_opt.seed}'
|
||||
elif opt.with_variations is not None:
|
||||
normalized_prompt = PromptFormatter(
|
||||
t2i, opt).normalize_prompt()
|
||||
# use the original seed - the per-iteration value is the last variation-seed
|
||||
metadata_prompt = f'{normalized_prompt} -S{opt.seed}'
|
||||
else:
|
||||
normalized_prompt = PromptFormatter(
|
||||
t2i, opt).normalize_prompt()
|
||||
metadata_prompt = f'{normalized_prompt} -S{seed}'
|
||||
path = file_writer.save_image_and_prompt_to_png(
|
||||
image, metadata_prompt, filename)
|
||||
if (not upscaled) or opt.save_original:
|
||||
# only append to results if we didn't overwrite an earlier output
|
||||
results.append([path, metadata_prompt])
|
||||
last_results.append([path, seed])
|
||||
|
||||
t2i.prompt2image(image_callback=image_writer, **vars(opt))
|
||||
|
||||
if do_grid and len(grid_images) > 0:
|
||||
grid_img = make_grid(list(grid_images.values()))
|
||||
grid_seeds = list(grid_images.keys())
|
||||
first_seed = last_results[0][1]
|
||||
filename = f'{prefix}.{first_seed}.png'
|
||||
# TODO better metadata for grid images
|
||||
normalized_prompt = PromptFormatter(
|
||||
t2i, opt).normalize_prompt()
|
||||
metadata_prompt = f'{normalized_prompt} -S{first_seed} --grid -n{len(grid_images)} # {grid_seeds}'
|
||||
path = file_writer.save_image_and_prompt_to_png(
|
||||
grid_img, metadata_prompt, filename
|
||||
)
|
||||
results = [[path, metadata_prompt]]
|
||||
|
||||
except AssertionError as e:
|
||||
print(e)
|
||||
continue
|
||||
|
||||
except OSError as e:
|
||||
print(e)
|
||||
continue
|
||||
|
||||
print('Outputs:')
|
||||
log_path = os.path.join(current_outdir, 'dream_log.txt')
|
||||
write_log_message(results, log_path)
|
||||
print()
|
||||
|
||||
print('goodbye!')
|
||||
|
||||
|
||||
def get_next_command(infile=None) -> str: # command string
|
||||
if infile is None:
|
||||
command = input('dream> ')
|
||||
else:
|
||||
command = infile.readline()
|
||||
if not command:
|
||||
raise EOFError
|
||||
else:
|
||||
command = command.strip()
|
||||
print(f'#{command}')
|
||||
return command
|
||||
|
||||
|
||||
def dream_server_loop(t2i, host, port, outdir):
|
||||
print('\n* --web was specified, starting web server...')
|
||||
# Change working directory to the stable-diffusion directory
|
||||
os.chdir(
|
||||
os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
|
||||
)
|
||||
|
||||
# Start server
|
||||
DreamServer.model = t2i
|
||||
DreamServer.outdir = outdir
|
||||
dream_server = ThreadingDreamServer((host, port))
|
||||
print(">> Started Stable Diffusion dream server!")
|
||||
if host == '0.0.0.0':
|
||||
print(
|
||||
f"Point your browser at http://localhost:{port} or use the host's DNS name or IP address.")
|
||||
else:
|
||||
print(">> Default host address now 127.0.0.1 (localhost). Use --host 0.0.0.0 to bind any address.")
|
||||
print(f">> Point your browser at http://{host}:{port}.")
|
||||
|
||||
try:
|
||||
dream_server.serve_forever()
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
dream_server.server_close()
|
||||
|
||||
|
||||
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='')
|
||||
|
||||
|
||||
with open(log_path, 'a', encoding='utf-8') as file:
|
||||
file.writelines(log_lines)
|
||||
|
||||
|
||||
SAMPLER_CHOICES = [
|
||||
'ddim',
|
||||
'k_dpm_2_a',
|
||||
'k_dpm_2',
|
||||
'k_euler_a',
|
||||
'k_euler',
|
||||
'k_heun',
|
||||
'k_lms',
|
||||
'plms',
|
||||
]
|
||||
|
||||
|
||||
def create_argv_parser():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="""Generate images using Stable Diffusion.
|
||||
Use --web to launch the web interface.
|
||||
Use --from_file to load prompts from a file path or standard input ("-").
|
||||
Otherwise you will be dropped into an interactive command prompt (type -h for help.)
|
||||
Other command-line arguments are defaults that can usually be overridden
|
||||
prompt the command prompt.
|
||||
"""
|
||||
)
|
||||
parser.add_argument(
|
||||
'--laion400m',
|
||||
'--latent_diffusion',
|
||||
'-l',
|
||||
dest='laion400m',
|
||||
action='store_true',
|
||||
help='Fallback to the latent diffusion (laion400m) weights and config',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--from_file',
|
||||
dest='infile',
|
||||
type=str,
|
||||
help='If specified, load prompts from this file',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-n',
|
||||
'--iterations',
|
||||
type=int,
|
||||
default=1,
|
||||
help='Number of images to generate',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-F',
|
||||
'--full_precision',
|
||||
dest='full_precision',
|
||||
action='store_true',
|
||||
help='Use more memory-intensive full precision math for calculations',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-g',
|
||||
'--grid',
|
||||
action='store_true',
|
||||
help='Generate a grid instead of individual images',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-A',
|
||||
'-m',
|
||||
'--sampler',
|
||||
dest='sampler_name',
|
||||
choices=SAMPLER_CHOICES,
|
||||
metavar='SAMPLER_NAME',
|
||||
default='k_lms',
|
||||
help=f'Set the initial sampler. Default: k_lms. Supported samplers: {", ".join(SAMPLER_CHOICES)}',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--outdir',
|
||||
'-o',
|
||||
type=str,
|
||||
default='outputs/img-samples',
|
||||
help='Directory to save generated images and a log of prompts and seeds. Default: outputs/img-samples',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--seamless',
|
||||
action='store_true',
|
||||
help='Change the model to seamless tiling (circular) mode',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--embedding_path',
|
||||
type=str,
|
||||
help='Path to a pre-trained embedding manager checkpoint - can only be set on command line',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--prompt_as_dir',
|
||||
'-p',
|
||||
action='store_true',
|
||||
help='Place images in subdirectories named after the prompt.',
|
||||
)
|
||||
# GFPGAN related args
|
||||
parser.add_argument(
|
||||
'--gfpgan_bg_upsampler',
|
||||
type=str,
|
||||
default='realesrgan',
|
||||
help='Background upsampler. Default: realesrgan. Options: realesrgan, none.',
|
||||
|
||||
)
|
||||
parser.add_argument(
|
||||
'--gfpgan_bg_tile',
|
||||
type=int,
|
||||
default=400,
|
||||
help='Tile size for background sampler, 0 for no tile during testing. Default: 400.',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--gfpgan_model_path',
|
||||
type=str,
|
||||
default='experiments/pretrained_models/GFPGANv1.3.pth',
|
||||
help='Indicates the path to the GFPGAN model, relative to --gfpgan_dir.',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--gfpgan_dir',
|
||||
type=str,
|
||||
default='./src/gfpgan',
|
||||
help='Indicates the directory containing the GFPGAN code.',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--web',
|
||||
dest='web',
|
||||
action='store_true',
|
||||
help='Start in web server mode.',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--host',
|
||||
type=str,
|
||||
default='127.0.0.1',
|
||||
help='Web server: Host or IP to listen on. Set to 0.0.0.0 to accept traffic from other devices on your network.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--port',
|
||||
type=int,
|
||||
default='9090',
|
||||
help='Web server: Port to listen on'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--weights',
|
||||
default='model',
|
||||
help='Indicates the Stable Diffusion model to use.',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--device',
|
||||
'-d',
|
||||
type=str,
|
||||
default='cuda',
|
||||
help="device to run stable diffusion on. defaults to cuda `torch.cuda.current_device()` if available"
|
||||
)
|
||||
parser.add_argument(
|
||||
'--model',
|
||||
default='stable-diffusion-1.4',
|
||||
help='Indicates which diffusion model to load. (currently "stable-diffusion-1.4" (default) or "laion400m")',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--config',
|
||||
default='configs/models.yaml',
|
||||
help='Path to configuration file for alternate models.',
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def create_cmd_parser():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Example: dream> a fantastic alien landscape -W1024 -H960 -s100 -n12'
|
||||
)
|
||||
parser.add_argument('prompt')
|
||||
parser.add_argument('-s', '--steps', type=int, help='Number of steps')
|
||||
parser.add_argument(
|
||||
'-S',
|
||||
'--seed',
|
||||
type=int,
|
||||
help='Image seed; a +ve integer, or use -1 for the previous seed, -2 for the one before that, etc',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-n',
|
||||
'--iterations',
|
||||
type=int,
|
||||
default=1,
|
||||
help='Number of samplings to perform (slower, but will provide seeds for individual images)',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-W', '--width', type=int, help='Image width, multiple of 64'
|
||||
)
|
||||
parser.add_argument(
|
||||
'-H', '--height', type=int, help='Image height, multiple of 64'
|
||||
)
|
||||
parser.add_argument(
|
||||
'-C',
|
||||
'--cfg_scale',
|
||||
default=7.5,
|
||||
type=float,
|
||||
help='Classifier free guidance (CFG) scale - higher numbers cause generator to "try" harder.',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-g', '--grid', action='store_true', help='generate a grid'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--outdir',
|
||||
'-o',
|
||||
type=str,
|
||||
default=None,
|
||||
help='Directory to save generated images and a log of prompts and seeds',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--seamless',
|
||||
action='store_true',
|
||||
help='Change the model to seamless tiling (circular) mode',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-i',
|
||||
'--individual',
|
||||
action='store_true',
|
||||
help='Generate individual files (default)',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-I',
|
||||
'--init_img',
|
||||
type=str,
|
||||
help='Path to input image for img2img mode (supersedes width and height)',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-M',
|
||||
'--init_mask',
|
||||
type=str,
|
||||
help='Path to input mask for inpainting mode (supersedes width and height)',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-T',
|
||||
'-fit',
|
||||
'--fit',
|
||||
action='store_true',
|
||||
help='If specified, will resize the input image to fit within the dimensions of width x height (512x512 default)',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-f',
|
||||
'--strength',
|
||||
default=0.75,
|
||||
type=float,
|
||||
help='Strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-G',
|
||||
'--gfpgan_strength',
|
||||
default=0,
|
||||
type=float,
|
||||
help='The strength at which to apply the GFPGAN model to the result, in order to improve faces.',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-U',
|
||||
'--upscale',
|
||||
nargs='+',
|
||||
default=None,
|
||||
type=float,
|
||||
help='Scale factor (2, 4) for upscaling followed by upscaling strength (0-1.0). If strength not specified, defaults to 0.75'
|
||||
)
|
||||
parser.add_argument(
|
||||
'-save_orig',
|
||||
'--save_original',
|
||||
action='store_true',
|
||||
help='Save original. Use it when upscaling to save both versions.',
|
||||
)
|
||||
# variants is going to be superseded by a generalized "prompt-morph" function
|
||||
# parser.add_argument('-v','--variants',type=int,help="in img2img mode, the first generated image will get passed back to img2img to generate the requested number of variants")
|
||||
parser.add_argument(
|
||||
'-x',
|
||||
'--skip_normalize',
|
||||
action='store_true',
|
||||
help='Skip subprompt weight normalization',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-A',
|
||||
'-m',
|
||||
'--sampler',
|
||||
dest='sampler_name',
|
||||
default=None,
|
||||
type=str,
|
||||
choices=SAMPLER_CHOICES,
|
||||
metavar='SAMPLER_NAME',
|
||||
help=f'Switch to a different sampler. Supported samplers: {", ".join(SAMPLER_CHOICES)}',
|
||||
)
|
||||
parser.add_argument(
|
||||
'-t',
|
||||
'--log_tokenization',
|
||||
action='store_true',
|
||||
help='shows how the prompt is split into tokens'
|
||||
)
|
||||
parser.add_argument(
|
||||
'-v',
|
||||
'--variation_amount',
|
||||
default=0.0,
|
||||
type=float,
|
||||
help='If > 0, generates variations on the initial seed instead of random seeds per iteration. Must be between 0 and 1. Higher values will be more different.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'-V',
|
||||
'--with_variations',
|
||||
default=None,
|
||||
type=str,
|
||||
help='list of variations to apply, in the format `seed:weight,seed:weight,...'
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
46
tests/legacy_tests.sh
Executable file
46
tests/legacy_tests.sh
Executable file
@ -0,0 +1,46 @@
|
||||
#! /usr/bin/env bash
|
||||
|
||||
# This file contains bunch of compatibility tests that ensures
|
||||
# that the API interface of `scripts/legacy-api.py` remains stable
|
||||
|
||||
set -e
|
||||
|
||||
OUTDIR=$(mktemp -d)
|
||||
|
||||
echo "Using directory $OUTDIR"
|
||||
|
||||
# Start API
|
||||
python -u scripts/legacy_api.py --web --host=localhost --port=3333 --outdir=$OUTDIR &> $OUTDIR/sd.log &
|
||||
APP_PID=$!
|
||||
|
||||
echo "Wait for server to startup"
|
||||
|
||||
tail -f -n0 $OUTDIR/sd.log | grep -qe "Point your browser at"
|
||||
|
||||
echo "Started, continuing"
|
||||
|
||||
if [ $? == 1 ]; then
|
||||
echo "Search terminated without finding the pattern"
|
||||
fi
|
||||
|
||||
# Generate image
|
||||
RESULT=$(curl -v -X POST -d '{"index":0,"variation_amount":0,"with_variations":"","steps":25,"width":512,"seed":"1337","prompt":"A cat wearing a hat","strength":0.5,"initimg":null,"cfg_scale":2,"iterations":1,"upscale_level":0,"upscale_strength":0,"sampler_name":"k_euler","height":512}' localhost:3333 | grep result)
|
||||
|
||||
# Test 01 - Image contents
|
||||
FILENAME=$(echo $RESULT | jq -r .url)
|
||||
|
||||
ACTUAL_CHECKSUM=$(sha256sum $FILENAME)
|
||||
EXPECTED_CHECKSUM="a77799226a4dfc62a1674498e575c775da042959a4b90b13e26f666c302f079f"
|
||||
|
||||
if [ "$ACTUAL_CHECKSUM" != "$EXPECTED_CHECKSUM" ]; then
|
||||
echo "Expected hash $EXPECTED_CHECKSUM but got hash $ACTUAL_CHECKSUM"
|
||||
kill $APP_PID
|
||||
# rm -r $OUTDIR
|
||||
exit 33
|
||||
fi
|
||||
|
||||
# Assert output
|
||||
|
||||
# Cleanup
|
||||
kill $APP_PID
|
||||
# rm -r $OUTDIR
|
||||
Loading…
Reference in New Issue
Block a user