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 development

Browse Source
This commit is contained in:
Lincoln Stein 2022-09-28 17:35:30 -04:00 committed by GitHub
commit b947290801
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
15 changed files with 1480 additions and 256 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
README.md
View File

@ -25,7 +25,7 @@ _This repository was formally known as lstein/stable-diffusion_
[CI checks on dev link]: https://github.com/invoke-ai/InvokeAI/actions?query=branch%3Adevelopment [CI checks on dev link]: https://github.com/invoke-ai/InvokeAI/actions?query=branch%3Adevelopment
[CI checks on main badge]: https://flat.badgen.net/github/checks/invoke-ai/InvokeAI/main?label=CI%20status%20on%20main&cache=900&icon=github [CI checks on main badge]: https://flat.badgen.net/github/checks/invoke-ai/InvokeAI/main?label=CI%20status%20on%20main&cache=900&icon=github
[CI checks on main link]: https://github.com/invoke-ai/InvokeAI/actions/workflows/test-dream-conda.yml [CI checks on main link]: https://github.com/invoke-ai/InvokeAI/actions/workflows/test-dream-conda.yml
[discord badge]: https://flat.badgen.net/discord/members/htRgbc7e?icon=discord [discord badge]: https://flat.badgen.net/discord/members/ZmtBAhwWhy?icon=discord
[discord link]: https://discord.gg/ZmtBAhwWhy [discord link]: https://discord.gg/ZmtBAhwWhy
[github forks badge]: https://flat.badgen.net/github/forks/invoke-ai/InvokeAI?icon=github [github forks badge]: https://flat.badgen.net/github/forks/invoke-ai/InvokeAI?icon=github
[github forks link]: https://useful-forks.github.io/?repo=invoke-ai%2FInvokeAI [github forks link]: https://useful-forks.github.io/?repo=invoke-ai%2FInvokeAI

880
backend/invoke_ai_web_server.py Normal file
View File

@ -0,0 +1,880 @@
import eventlet
import glob
import os
import shutil
import mimetypes
from flask import Flask, redirect, send_from_directory
from flask_socketio import SocketIO
from PIL import Image
from uuid import uuid4
from threading import Event
from ldm.dream.args import Args, APP_ID, APP_VERSION, calculate_init_img_hash
from ldm.dream.pngwriter import PngWriter, retrieve_metadata
from ldm.dream.conditioning import split_weighted_subprompts
from backend.modules.parameters import parameters_to_command
# Loading Arguments
opt = Args()
args = opt.parse_args()
class InvokeAIWebServer:
def __init__(self, generate, gfpgan, codeformer, esrgan) -> None:
self.host = args.host
self.port = args.port
self.generate = generate
self.gfpgan = gfpgan
self.codeformer = codeformer
self.esrgan = esrgan
self.canceled = Event()
def run(self):
self.setup_app()
self.setup_flask()
def setup_flask(self):
# Fix missing mimetypes on Windows
mimetypes.add_type("application/javascript", ".js")
mimetypes.add_type("text/css", ".css")
# Socket IO
logger = True if args.web_verbose else False
engineio_logger = True if args.web_verbose else False
max_http_buffer_size = 10000000
# CORS Allowed Setup
cors_allowed_origins = ['http://127.0.0.1:5173', 'http://localhost:5173']
additional_allowed_origins = (
opt.cors if opt.cors else []
) # additional CORS allowed origins
if self.host == '127.0.0.1':
cors_allowed_origins.extend(
[
f'http://{self.host}:{self.port}',
f'http://localhost:{self.port}',
]
)
cors_allowed_origins = (
cors_allowed_origins + additional_allowed_origins
)
self.app = Flask(
__name__, static_url_path='', static_folder='../frontend/dist/'
)
self.socketio = SocketIO(
self.app,
logger=logger,
engineio_logger=engineio_logger,
max_http_buffer_size=max_http_buffer_size,
cors_allowed_origins=cors_allowed_origins,
ping_interval=(50, 50),
ping_timeout=60,
)
# Outputs Route
self.app.config['OUTPUTS_FOLDER'] = os.path.abspath(args.outdir)
@self.app.route('/outputs/<path:file_path>')
def outputs(file_path):
return send_from_directory(
self.app.config['OUTPUTS_FOLDER'], file_path
)
# Base Route
@self.app.route('/')
def serve():
if args.web_develop:
return redirect('http://127.0.0.1:5173')
else:
return send_from_directory(
self.app.static_folder, 'index.html'
)
self.load_socketio_listeners(self.socketio)
print('>> Started Invoke AI Web Server!')
if self.host == '0.0.0.0':
print(
f"Point your browser at http://localhost:{self.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://{self.host}:{self.port}')
self.socketio.run(app=self.app, host=self.host, port=self.port)
def setup_app(self):
self.result_url = 'outputs/'
self.init_image_url = 'outputs/init-images/'
self.mask_image_url = 'outputs/mask-images/'
self.intermediate_url = 'outputs/intermediates/'
# location for "finished" images
self.result_path = args.outdir
# temporary path for intermediates
self.intermediate_path = os.path.join(
self.result_path, 'intermediates/'
)
# path for user-uploaded init images and masks
self.init_image_path = os.path.join(self.result_path, 'init-images/')
self.mask_image_path = os.path.join(self.result_path, 'mask-images/')
# txt log
self.log_path = os.path.join(self.result_path, 'dream_log.txt')
# make all output paths
[
os.makedirs(path, exist_ok=True)
for path in [
self.result_path,
self.intermediate_path,
self.init_image_path,
self.mask_image_path,
]
]
def load_socketio_listeners(self, socketio):
@socketio.on('requestSystemConfig')
def handle_request_capabilities():
print(f'>> System config requested')
config = self.get_system_config()
socketio.emit('systemConfig', config)
@socketio.on('requestImages')
def handle_request_images(page=1, offset=0, last_mtime=None):
chunk_size = 50
if last_mtime:
print(f'>> Latest images requested')
else:
print(
f'>> Page {page} of images requested (page size {chunk_size} offset {offset})'
)
paths = glob.glob(os.path.join(self.result_path, '*.png'))
sorted_paths = sorted(
paths, key=lambda x: os.path.getmtime(x), reverse=True
)
if last_mtime:
image_paths = filter(
lambda x: os.path.getmtime(x) > last_mtime, sorted_paths
)
else:
image_paths = sorted_paths[
slice(
chunk_size * (page - 1) + offset,
chunk_size * page + offset,
)
]
page = page + 1
image_array = []
for path in image_paths:
metadata = retrieve_metadata(path)
image_array.append(
{
'url': self.get_url_from_image_path(path),
'mtime': os.path.getmtime(path),
'metadata': metadata['sd-metadata'],
}
)
socketio.emit(
'galleryImages',
{
'images': image_array,
'nextPage': page,
'offset': offset,
'onlyNewImages': True if last_mtime else False,
},
)
@socketio.on('generateImage')
def handle_generate_image_event(
generation_parameters, esrgan_parameters, gfpgan_parameters
):
print(
f'>> Image generation requested: {generation_parameters}\nESRGAN parameters: {esrgan_parameters}\nGFPGAN parameters: {gfpgan_parameters}'
)
self.generate_images(
generation_parameters, esrgan_parameters, gfpgan_parameters
)
@socketio.on('runESRGAN')
def handle_run_esrgan_event(original_image, esrgan_parameters):
print(
f'>> ESRGAN upscale requested for "{original_image["url"]}": {esrgan_parameters}'
)
progress = {
'currentStep': 1,
'totalSteps': 1,
'currentIteration': 1,
'totalIterations': 1,
'currentStatus': 'Preparing',
'isProcessing': True,
'currentStatusHasSteps': False,
}
socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
original_image_path = self.get_image_path_from_url(original_image['url'])
# os.path.join(self.result_path, os.path.basename(original_image['url']))
image = Image.open(original_image_path)
seed = (
original_image['metadata']['seed']
if 'seed' in original_image['metadata']
else 'unknown_seed'
)
progress['currentStatus'] = 'Upscaling'
socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
image = self.esrgan.process(
image=image,
upsampler_scale=esrgan_parameters['upscale'][0],
strength=esrgan_parameters['upscale'][1],
seed=seed,
)
progress['currentStatus'] = 'Saving image'
socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
esrgan_parameters['seed'] = seed
metadata = self.parameters_to_post_processed_image_metadata(
parameters=esrgan_parameters,
original_image_path=original_image_path,
type='esrgan',
)
command = parameters_to_command(esrgan_parameters)
path = self.save_image(
image,
command,
metadata,
self.result_path,
postprocessing='esrgan',
)
self.write_log_message(
f'[Upscaled] "{original_image_path}" > "{path}": {command}'
)
progress['currentStatus'] = 'Finished'
progress['currentStep'] = 0
progress['totalSteps'] = 0
progress['currentIteration'] = 0
progress['totalIterations'] = 0
progress['isProcessing'] = False
socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
socketio.emit(
'esrganResult',
{
'url': self.get_url_from_image_path(path),
'mtime': os.path.getmtime(path),
'metadata': metadata,
},
)
@socketio.on('runGFPGAN')
def handle_run_gfpgan_event(original_image, gfpgan_parameters):
print(
f'>> GFPGAN face fix requested for "{original_image["url"]}": {gfpgan_parameters}'
)
progress = {
'currentStep': 1,
'totalSteps': 1,
'currentIteration': 1,
'totalIterations': 1,
'currentStatus': 'Preparing',
'isProcessing': True,
'currentStatusHasSteps': False,
}
socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
original_image_path = self.get_image_path_from_url(original_image['url'])
image = Image.open(original_image_path)
seed = (
original_image['metadata']['seed']
if 'seed' in original_image['metadata']
else 'unknown_seed'
)
progress['currentStatus'] = 'Fixing faces'
socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
image = self.gfpgan.process(
image=image,
strength=gfpgan_parameters['gfpgan_strength'],
seed=seed,
)
progress['currentStatus'] = 'Saving image'
socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
gfpgan_parameters['seed'] = seed
metadata = self.parameters_to_post_processed_image_metadata(
parameters=gfpgan_parameters,
original_image_path=original_image_path,
type='gfpgan',
)
command = parameters_to_command(gfpgan_parameters)
path = self.save_image(
image,
command,
metadata,
self.result_path,
postprocessing='gfpgan',
)
self.write_log_message(
f'[Fixed faces] "{original_image_path}" > "{path}": {command}'
)
progress['currentStatus'] = 'Finished'
progress['currentStep'] = 0
progress['totalSteps'] = 0
progress['currentIteration'] = 0
progress['totalIterations'] = 0
progress['isProcessing'] = False
socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
socketio.emit(
'gfpganResult',
{
'url': self.get_url_from_image_path(path),
'mtime': os.path.getmtime(path),
'metadata': metadata,
},
)
@socketio.on('cancel')
def handle_cancel():
print(f'>> Cancel processing requested')
self.canceled.set()
socketio.emit('processingCanceled')
# TODO: I think this needs a safety mechanism.
@socketio.on('deleteImage')
def handle_delete_image(path, uuid):
print(f'>> Delete requested "{path}"')
from send2trash import send2trash
path = self.get_image_path_from_url(path)
send2trash(path)
socketio.emit('imageDeleted', {'url': path, 'uuid': uuid})
# TODO: I think this needs a safety mechanism.
@socketio.on('uploadInitialImage')
def handle_upload_initial_image(bytes, name):
print(f'>> Init image upload requested "{name}"')
uuid = uuid4().hex
split = os.path.splitext(name)
name = f'{split[0]}.{uuid}{split[1]}'
file_path = os.path.join(self.init_image_path, name)
os.makedirs(os.path.dirname(file_path), exist_ok=True)
newFile = open(file_path, 'wb')
newFile.write(bytes)
socketio.emit(
'initialImageUploaded', {'url': self.get_url_from_image_path(file_path), 'uuid': ''}
)
# TODO: I think this needs a safety mechanism.
@socketio.on('uploadMaskImage')
def handle_upload_mask_image(bytes, name):
print(f'>> Mask image upload requested "{name}"')
uuid = uuid4().hex
split = os.path.splitext(name)
name = f'{split[0]}.{uuid}{split[1]}'
file_path = os.path.join(self.mask_image_path, name)
os.makedirs(os.path.dirname(file_path), exist_ok=True)
newFile = open(file_path, 'wb')
newFile.write(bytes)
socketio.emit('maskImageUploaded', {'url': self.get_url_from_image_path(file_path), 'uuid': ''})
# App Functions
def get_system_config(self):
return {
'model': 'stable diffusion',
'model_id': args.model,
'model_hash': self.generate.model_hash,
'app_id': APP_ID,
'app_version': APP_VERSION,
}
def generate_images(
self, generation_parameters, esrgan_parameters, gfpgan_parameters
):
self.canceled.clear()
step_index = 1
prior_variations = (
generation_parameters['with_variations']
if 'with_variations' in generation_parameters
else []
)
"""
TODO: RE-IMPLEMENT THE COMMENTED-OUT CODE
If a result image is used as an init image, and then deleted, we will want to be
able to use it as an init image in the future. Need to copy it.
If the init/mask image doesn't exist in the init_image_path/mask_image_path,
make a unique filename for it and copy it there.
"""
# if 'init_img' in generation_parameters:
# filename = os.path.basename(generation_parameters['init_img'])
# abs_init_image_path = os.path.join(self.init_image_path, filename)
# if not os.path.exists(
# abs_init_image_path
# ):
# unique_filename = self.make_unique_init_image_filename(
# filename
# )
# new_path = os.path.join(self.init_image_path, unique_filename)
# shutil.copy(abs_init_image_path, new_path)
# generation_parameters['init_img'] = os.path.abspath(new_path)
# else:
# generation_parameters['init_img'] = os.path.abspath(os.path.join(self.init_image_path, filename))
# if 'init_mask' in generation_parameters:
# filename = os.path.basename(generation_parameters['init_mask'])
# if not os.path.exists(
# os.path.join(self.mask_image_path, filename)
# ):
# unique_filename = self.make_unique_init_image_filename(
# filename
# )
# new_path = os.path.join(
# self.init_image_path, unique_filename
# )
# shutil.copy(generation_parameters['init_img'], new_path)
# generation_parameters['init_mask'] = os.path.abspath(new_path)
# else:
# generation_parameters['init_mas'] = os.path.abspath(os.path.join(self.mask_image_path, filename))
# We need to give absolute paths to the generator, stash the URLs for later
init_img_url = None;
mask_img_url = None;
if 'init_img' in generation_parameters:
init_img_url = generation_parameters['init_img']
generation_parameters['init_img'] = self.get_image_path_from_url(generation_parameters['init_img'])
if 'init_mask' in generation_parameters:
mask_img_url = generation_parameters['init_mask']
generation_parameters['init_mask'] = self.get_image_path_from_url(generation_parameters['init_mask'])
totalSteps = self.calculate_real_steps(
steps=generation_parameters['steps'],
strength=generation_parameters['strength']
if 'strength' in generation_parameters
else None,
has_init_image='init_img' in generation_parameters,
)
progress = {
'currentStep': 1,
'totalSteps': totalSteps,
'currentIteration': 1,
'totalIterations': generation_parameters['iterations'],
'currentStatus': 'Preparing',
'isProcessing': True,
'currentStatusHasSteps': False,
}
self.socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
def image_progress(sample, step):
if self.canceled.is_set():
raise CanceledException
nonlocal step_index
nonlocal generation_parameters
nonlocal progress
progress['currentStep'] = step + 1
progress['currentStatus'] = 'Generating'
progress['currentStatusHasSteps'] = True
if (
generation_parameters['progress_images']
and step % 5 == 0
and step < generation_parameters['steps'] - 1
):
image = self.generate.sample_to_image(sample)
metadata = self.parameters_to_generated_image_metadata(generation_parameters)
command = parameters_to_command(generation_parameters)
path = self.save_image(image, command, metadata, self.intermediate_path, step_index=step_index, postprocessing=False)
step_index += 1
self.socketio.emit(
'intermediateResult',
{
'url': self.get_url_from_image_path(path),
'mtime': os.path.getmtime(path),
'metadata': metadata,
},
)
self.socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
def image_done(image, seed, first_seed):
nonlocal generation_parameters
nonlocal esrgan_parameters
nonlocal gfpgan_parameters
nonlocal progress
step_index = 1
nonlocal prior_variations
progress['currentStatus'] = 'Generation complete'
self.socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
all_parameters = generation_parameters
postprocessing = False
if (
'variation_amount' in all_parameters
and all_parameters['variation_amount'] > 0
):
first_seed = first_seed or seed
this_variation = [[seed, all_parameters['variation_amount']]]
all_parameters['with_variations'] = (
prior_variations + this_variation
)
all_parameters['seed'] = first_seed
elif 'with_variations' in all_parameters:
all_parameters['seed'] = first_seed
else:
all_parameters['seed'] = seed
if esrgan_parameters:
progress['currentStatus'] = 'Upscaling'
progress['currentStatusHasSteps'] = False
self.socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
image = self.esrgan.process(
image=image,
upsampler_scale=esrgan_parameters['level'],
strength=esrgan_parameters['strength'],
seed=seed,
)
postprocessing = True
all_parameters['upscale'] = [
esrgan_parameters['level'],
esrgan_parameters['strength'],
]
if gfpgan_parameters:
progress['currentStatus'] = 'Fixing faces'
progress['currentStatusHasSteps'] = False
self.socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
image = self.gfpgan.process(
image=image,
strength=gfpgan_parameters['strength'],
seed=seed,
)
postprocessing = True
all_parameters['gfpgan_strength'] = gfpgan_parameters[
'strength'
]
progress['currentStatus'] = 'Saving image'
self.socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
# restore the stashed URLS and discard the paths, we are about to send the result to client
if 'init_img' in all_parameters:
all_parameters['init_img'] = init_img_url
if 'init_mask' in all_parameters:
all_parameters['init_mask'] = mask_img_url
metadata = self.parameters_to_generated_image_metadata(
all_parameters
)
command = parameters_to_command(all_parameters)
path = self.save_image(
image,
command,
metadata,
self.result_path,
postprocessing=postprocessing,
)
print(f'>> Image generated: "{path}"')
self.write_log_message(f'[Generated] "{path}": {command}')
if progress['totalIterations'] > progress['currentIteration']:
progress['currentStep'] = 1
progress['currentIteration'] += 1
progress['currentStatus'] = 'Iteration finished'
progress['currentStatusHasSteps'] = False
else:
progress['currentStep'] = 0
progress['totalSteps'] = 0
progress['currentIteration'] = 0
progress['totalIterations'] = 0
progress['currentStatus'] = 'Finished'
progress['isProcessing'] = False
self.socketio.emit('progressUpdate', progress)
eventlet.sleep(0)
self.socketio.emit(
'generationResult',
{
'url': self.get_url_from_image_path(path),
'mtime': os.path.getmtime(path),
'metadata': metadata,
},
)
eventlet.sleep(0)
try:
self.generate.prompt2image(
**generation_parameters,
step_callback=image_progress,
image_callback=image_done,
)
except KeyboardInterrupt:
raise
except CanceledException:
pass
except Exception as e:
self.socketio.emit('error', {'message': (str(e))})
print('\n')
import traceback
traceback.print_exc()
print('\n')
def parameters_to_generated_image_metadata(self, parameters):
# top-level metadata minus `image` or `images`
metadata = self.get_system_config()
# remove any image keys not mentioned in RFC #266
rfc266_img_fields = [
'type',
'postprocessing',
'sampler',
'prompt',
'seed',
'variations',
'steps',
'cfg_scale',
'step_number',
'width',
'height',
'extra',
'seamless',
]
rfc_dict = {}
for item in parameters.items():
key, value = item
if key in rfc266_img_fields:
rfc_dict[key] = value
postprocessing = []
# 'postprocessing' is either null or an
if 'gfpgan_strength' in parameters:
postprocessing.append(
{
'type': 'gfpgan',
'strength': float(parameters['gfpgan_strength']),
}
)
if 'upscale' in parameters:
postprocessing.append(
{
'type': 'esrgan',
'scale': int(parameters['upscale'][0]),
'strength': float(parameters['upscale'][1]),
}
)
rfc_dict['postprocessing'] = (
postprocessing if len(postprocessing) > 0 else None
)
# semantic drift
rfc_dict['sampler'] = parameters['sampler_name']
# display weighted subprompts (liable to change)
subprompts = split_weighted_subprompts(parameters['prompt'])
subprompts = [{'prompt': x[0], 'weight': x[1]} for x in subprompts]
rfc_dict['prompt'] = subprompts
# 'variations' should always exist and be an array, empty or consisting of {'seed': seed, 'weight': weight} pairs
variations = []
if 'with_variations' in parameters:
variations = [
{'seed': x[0], 'weight': x[1]}
for x in parameters['with_variations']
]
rfc_dict['variations'] = variations
if 'init_img' in parameters:
rfc_dict['type'] = 'img2img'
rfc_dict['strength'] = parameters['strength']
rfc_dict['fit'] = parameters['fit'] # TODO: Noncompliant
rfc_dict['orig_hash'] = calculate_init_img_hash(self.get_image_path_from_url(parameters['init_img']))
rfc_dict['init_image_path'] = parameters[
'init_img'
] # TODO: Noncompliant
rfc_dict[
'sampler'
] = 'ddim' # TODO: FIX ME WHEN IMG2IMG SUPPORTS ALL SAMPLERS
if 'init_mask' in parameters:
rfc_dict['mask_hash'] = calculate_init_img_hash(self.get_image_path_from_url(parameters['init_mask'])) # TODO: Noncompliant
rfc_dict['mask_image_path'] = parameters[
'init_mask'
] # TODO: Noncompliant
else:
rfc_dict['type'] = 'txt2img'
metadata['image'] = rfc_dict
return metadata
def parameters_to_post_processed_image_metadata(
self, parameters, original_image_path, type
):
# top-level metadata minus `image` or `images`
metadata = self.get_system_config()
orig_hash = calculate_init_img_hash(self.get_image_path_from_url(original_image_path))
image = {'orig_path': original_image_path, 'orig_hash': orig_hash}
if type == 'esrgan':
image['type'] = 'esrgan'
image['scale'] = parameters['upscale'][0]
image['strength'] = parameters['upscale'][1]
elif type == 'gfpgan':
image['type'] = 'gfpgan'
image['strength'] = parameters['gfpgan_strength']
else:
raise TypeError(f'Invalid type: {type}')
metadata['image'] = image
return metadata
def save_image(
self,
image,
command,
metadata,
output_dir,
step_index=None,
postprocessing=False,
):
pngwriter = PngWriter(output_dir)
prefix = pngwriter.unique_prefix()
seed = 'unknown_seed'
if 'image' in metadata:
if 'seed' in metadata['image']:
seed = metadata['image']['seed']
filename = f'{prefix}.{seed}'
if step_index:
filename += f'.{step_index}'
if postprocessing:
filename += f'.postprocessed'
filename += '.png'
path = pngwriter.save_image_and_prompt_to_png(
image=image, dream_prompt=command, metadata=metadata, name=filename
)
return os.path.abspath(path)
def make_unique_init_image_filename(self, name):
uuid = uuid4().hex
split = os.path.splitext(name)
name = f'{split[0]}.{uuid}{split[1]}'
return name
def calculate_real_steps(self, steps, strength, has_init_image):
import math
return math.floor(strength * steps) if has_init_image else steps
def write_log_message(self, message):
"""Logs the filename and parameters used to generate or process that image to log file"""
message = f'{message}\n'
with open(self.log_path, 'a', encoding='utf-8') as file:
file.writelines(message)
def get_image_path_from_url(self, url):
"""Given a url to an image used by the client, returns the absolute file path to that image"""
if 'init-images' in url:
return os.path.abspath(os.path.join(self.init_image_path, os.path.basename(url)))
elif 'mask-images' in url:
return os.path.abspath(os.path.join(self.mask_image_path, os.path.basename(url)))
elif 'intermediates' in url:
return os.path.abspath(os.path.join(self.intermediate_path, os.path.basename(url)))
else:
return os.path.abspath(os.path.join(self.result_path, os.path.basename(url)))
def get_url_from_image_path(self, path):
"""Given an absolute file path to an image, returns the URL that the client can use to load the image"""
if 'init-images' in path:
return os.path.join(self.init_image_url, os.path.basename(path))
elif 'mask-images' in path:
return os.path.join(self.mask_image_url, os.path.basename(path))
elif 'intermediates' in path:
return os.path.join(self.intermediate_url, os.path.basename(path))
else:
return os.path.join(self.result_url, os.path.basename(path))
class CanceledException(Exception):
pass

6
backend/modules/create_cmd_parser.py
View File

@ -45,5 +45,11 @@ def create_cmd_parser():
help=f'Set model precision. Defaults to auto selected based on device. Options: {", ".join(PRECISION_CHOICES)}', help=f'Set model precision. Defaults to auto selected based on device. Options: {", ".join(PRECISION_CHOICES)}',
default="auto", default="auto",
) )
parser.add_argument(
'--free_gpu_mem',
dest='free_gpu_mem',
action='store_true',
help='Force free gpu memory before final decoding',
)
return parser return parser

2
backend/modules/parameters.py
View File

@ -1,4 +1,4 @@
from modules.parse_seed_weights import parse_seed_weights from backend.modules.parse_seed_weights import parse_seed_weights
import argparse import argparse
SAMPLER_CHOICES = [ SAMPLER_CHOICES = [

2
backend/server.py
View File

@ -50,6 +50,7 @@ host = opt.host # Web & socket.io host
port = opt.port # Web & socket.io port port = opt.port # Web & socket.io port
verbose = opt.verbose # enables copious socket.io logging verbose = opt.verbose # enables copious socket.io logging
precision = opt.precision precision = opt.precision
free_gpu_mem = opt.free_gpu_mem
embedding_path = opt.embedding_path embedding_path = opt.embedding_path
additional_allowed_origins = ( additional_allowed_origins = (
opt.cors if opt.cors else [] opt.cors if opt.cors else []
@ -148,6 +149,7 @@ generate = Generate(
precision=precision, precision=precision,
embedding_path=embedding_path, embedding_path=embedding_path,
) )
generate.free_gpu_mem = free_gpu_mem
generate.load_model() generate.load_model()

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.| | --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 # Command-line editing and completion
If you are on a Macintosh or Linux machine, the command-line offers If you are on a Macintosh or Linux machine, the command-line offers

1
docs/help/TROUBLESHOOT.md
View File

@ -87,7 +87,6 @@ Usually this will be sufficient, but if you start to see errors about
missing or incorrect modules, use the command `pip install -e .` missing or incorrect modules, use the command `pip install -e .`
and/or `conda env update` (These commands won't break anything.) and/or `conda env update` (These commands won't break anything.)
`pip install -e .` and/or `pip install -e .` and/or
`conda env update -f environment.yaml` `conda env update -f environment.yaml`

19
docs/installation/INSTALL_MAC.md
View File

@ -118,16 +118,17 @@ ln -s "$PATH_TO_CKPT/sd-v1-4.ckpt" \
```bash ```bash
PIP_EXISTS_ACTION=w CONDA_SUBDIR=osx-arm64 \ PIP_EXISTS_ACTION=w CONDA_SUBDIR=osx-arm64 \
conda env create \ conda env create \
-f environment-mac.yaml \ -f environment-mac.yml \
&& conda activate ldm && conda activate ldm
``` ```
=== "Intel x86_64" === "Intel x86_64"
```bash ```bash
PIP_EXISTS_ACTION=w CONDA_SUBDIR=osx-64 \ PIP_EXISTS_ACTION=w CONDA_SUBDIR=osx-64 \
conda env create \ conda env create \
-f environment-mac.yaml \ -f environment-mac.yml \
&& conda activate ldm && conda activate ldm
``` ```
@ -147,16 +148,9 @@ python scripts/orig_scripts/txt2img.py \
--plms --plms
``` ```
## Notes Note, `export PIP_EXISTS_ACTION=w` is a precaution to fix `conda env
1. half-precision requires autocast which is unfortunately incompatible with the
implementation of pytorch on the M1 architecture. On Macs, --full-precision will
default to True.
2. `export PIP_EXISTS_ACTION=w` in the commands above, is a precaution to fix `conda env
create -f environment-mac.yml` never finishing in some situations. So create -f environment-mac.yml` never finishing in some situations. So
it isn't required but wont hurt. it isn't required but wont hurt.
--- ---
## Common problems ## Common problems
@ -196,7 +190,8 @@ conda install \
-n ldm -n ldm
``` ```
If it takes forever to run `conda env create -f environment-mac.yml` you could try to run:
If it takes forever to run `conda env create -f environment-mac.yml`, try this:
```bash ```bash
git clean -f git clean -f
@ -384,7 +379,7 @@ python scripts/preload_models.py
``` ```
This fork already includes a fix for this in This fork already includes a fix for this in
[environment-mac.yaml](https://github.com/invoke-ai/InvokeAI/blob/main/environment-mac.yml). [environment-mac.yml](https://github.com/invoke-ai/InvokeAI/blob/main/environment-mac.yml).
### "Could not build wheels for tokenizers" ### "Could not build wheels for tokenizers"

17
docs/installation/INSTALL_WINDOWS.md
View File

@ -39,7 +39,7 @@ in the wiki
4. Run the command: 4. Run the command:
```batch ```bash
git clone https://github.com/invoke-ai/InvokeAI.git git clone https://github.com/invoke-ai/InvokeAI.git
``` ```
@ -48,17 +48,16 @@ in the wiki
5. Enter the newly-created InvokeAI folder. From this step forward make sure that you are working in the InvokeAI directory! 5. Enter the newly-created InvokeAI folder. From this step forward make sure that you are working in the InvokeAI directory!
```batch ```
cd InvokeAI cd InvokeAI
``` ```
6. Run the following two commands: 6. Run the following two commands:
```batch ```
conda env create (step 6a) conda env create (step 6a)
conda activate ldm (step 6b) conda activate ldm (step 6b)
``` ```
This will install all python requirements and activate the "ldm" environment This will install all python requirements and activate the "ldm" environment
which sets PATH and other environment variables properly. which sets PATH and other environment variables properly.
@ -68,7 +67,7 @@ in the wiki
7. Run the command: 7. Run the command:
```batch ```bash
python scripts\preload_models.py python scripts\preload_models.py
``` ```
@ -90,7 +89,7 @@ in the wiki
Now run the following commands from **within the InvokeAI directory** to copy the weights file to the right place: Now run the following commands from **within the InvokeAI directory** to copy the weights file to the right place:
```batch ```
mkdir -p models\ldm\stable-diffusion-v1 mkdir -p models\ldm\stable-diffusion-v1
copy C:\path\to\sd-v1-4.ckpt models\ldm\stable-diffusion-v1\model.ckpt copy C:\path\to\sd-v1-4.ckpt models\ldm\stable-diffusion-v1\model.ckpt
``` ```
@ -100,7 +99,7 @@ you may instead create a shortcut to it from within `models\ldm\stable-diffusion
9. Start generating images! 9. Start generating images!
```batch ```bash
# for the pre-release weights # for the pre-release weights
python scripts\dream.py -l python scripts\dream.py -l
@ -117,14 +116,14 @@ you may instead create a shortcut to it from within `models\ldm\stable-diffusion
--- ---
## Updating to newer versions of the script This distribution is changing rapidly. If you used the `git clone` method (step 5) to download the InvokeAI directory, then to update to the latest and greatest version, launch the Anaconda window, enter `InvokeAI`, and type:
This distribution is changing rapidly. If you used the `git clone` method This distribution is changing rapidly. If you used the `git clone` method
(step 5) to download the stable-diffusion directory, then to update to the (step 5) to download the stable-diffusion directory, then to update to the
latest and greatest version, launch the Anaconda window, enter latest and greatest version, launch the Anaconda window, enter
`stable-diffusion`, and type: `stable-diffusion`, and type:
```batch ```bash
git pull git pull
conda env update conda env update
``` ```

2
docs/other/README-CompViz.md
View File

@ -40,7 +40,7 @@ A suitable [conda](https://conda.io/) environment named `ldm` can be created and
activated with: activated with:
``` ```
conda env create -f environment.yml conda env create
conda activate ldm conda activate ldm
``` ```

93
ldm/dream/args.py
View File

@ -81,13 +81,15 @@ with metadata_from_png():
""" """
import argparse import argparse
from argparse import Namespace from argparse import Namespace, RawTextHelpFormatter
import shlex import shlex
import json import json
import hashlib import hashlib
import os import os
import re
import copy import copy
import base64 import base64
import functools
import ldm.dream.pngwriter import ldm.dream.pngwriter
from ldm.dream.conditioning import split_weighted_subprompts from ldm.dream.conditioning import split_weighted_subprompts
@ -220,9 +222,15 @@ class Args(object):
# outpainting parameters # outpainting parameters
if a['out_direction']: if a['out_direction']:
switches.append(f'-D {" ".join([str(u) for u in 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']: if a['with_variations']:
formatted_variations = ','.join(f'{seed}:{weight}' for seed, weight in (a["with_variations"])) formatted_variations = ','.join(f'{seed}:{weight}' for seed, weight in (a["variations"]))
switches.append(f'-V {formatted_variations}') switches.append(f'-V {a["formatted_variations"]}')
if 'variations' in a:
switches.append(f'-V {a["variations"]}')
return ' '.join(switches) return ' '.join(switches)
def __getattribute__(self,name): def __getattribute__(self,name):
@ -421,6 +429,23 @@ class Args(object):
action='store_true', action='store_true',
help='Start in web server mode.', help='Start in web server mode.',
) )
web_server_group.add_argument(
'--web_develop',
dest='web_develop',
action='store_true',
help='Start in web server development mode.',
)
web_server_group.add_argument(
"--web_verbose",
action="store_true",
help="Enables verbose logging",
)
web_server_group.add_argument(
"--cors",
nargs="*",
type=str,
help="Additional allowed origins, comma-separated",
)
web_server_group.add_argument( web_server_group.add_argument(
'--host', '--host',
type=str, type=str,
@ -438,9 +463,24 @@ class Args(object):
# This creates the parser that processes commands on the dream> command line # This creates the parser that processes commands on the dream> command line
def _create_dream_cmd_parser(self): def _create_dream_cmd_parser(self):
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description=""" formatter_class=RawTextHelpFormatter,
Generate example: dream> a fantastic alien landscape -W576 -H512 -s60 -n4 description=
Postprocess example: dream> !pp 0000045.4829112.png -G1 -U4 -ft codeformer """
*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') render_group = parser.add_argument_group('General rendering')
@ -608,7 +648,7 @@ class Args(object):
'-embiggen', '-embiggen',
nargs='+', nargs='+',
type=float, 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, default=None,
) )
postprocessing_group.add_argument( postprocessing_group.add_argument(
@ -616,7 +656,7 @@ class Args(object):
'-embiggen_tiles', '-embiggen_tiles',
nargs='+', nargs='+',
type=int, 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, default=None,
) )
special_effects_group.add_argument( special_effects_group.add_argument(
@ -732,19 +772,29 @@ def metadata_dumps(opt,
return metadata return metadata
def metadata_from_png(png_file_path): @functools.lru_cache(maxsize=50)
def metadata_from_png(png_file_path) -> Args:
''' '''
Given the path to a PNG file created by dream.py, retrieves Given the path to a PNG file created by dream.py, retrieves
an Args object containing the image metadata an Args object containing the image metadata. Note that this
returns a single Args object, not multiple.
''' '''
meta = ldm.dream.pngwriter.retrieve_metadata(png_file_path) meta = ldm.dream.pngwriter.retrieve_metadata(png_file_path)
opts = metadata_loads(meta) if 'sd-metadata' in meta and len(meta['sd-metadata'])>0 :
return opts[0] return metadata_loads(meta)[0]
else:
return legacy_metadata_load(meta,png_file_path)
def metadata_loads(metadata): def dream_cmd_from_png(png_file_path):
opt = metadata_from_png(png_file_path)
return opt.dream_prompt_str()
def metadata_loads(metadata) -> list:
''' '''
Takes the dictionary corresponding to RFC266 (https://github.com/lstein/stable-diffusion/issues/266) Takes the dictionary corresponding to RFC266 (https://github.com/lstein/stable-diffusion/issues/266)
and returns a series of opt objects for each of the images described in the dictionary. and returns a series of opt objects for each of the images described in the dictionary. Note that this
returns a list, and not a single object. See metadata_from_png() for a more convenient function for
files that contain a single image.
''' '''
results = [] results = []
try: try:
@ -797,3 +847,18 @@ def sha256(path):
sha.update(data) sha.update(data)
return sha.hexdigest() return sha.hexdigest()
def legacy_metadata_load(meta,pathname) -> Args:
if 'Dream' in meta and len(meta['Dream']) > 0:
dream_prompt = meta['Dream']
opt = Args()
opt.parse_cmd(dream_prompt)
return opt
else: # if nothing else, we can get the seed
match = re.search('\d+\.(\d+)',pathname)
if match:
seed = match.groups()[0]
opt = Args()
opt.seed = seed
return opt
return None

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): def write_log_message(results, output_cntr):
"""logs to the terminal""" """logs to the terminal"""
log_lines = [f"{path}: {prompt}\n" for path, prompt in results] if len(results) == 0:
for l in log_lines:
output_cntr += 1
print(f"[{output_cntr}] {l}", end="")
return output_cntr return output_cntr
log_lines = [f"{path}: {prompt}\n" for path, prompt in results]
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): def write_log_files(results, log_path, file_types):
for file_type in file_types: for file_type in file_types:

312
ldm/dream/readline.py
View File

@ -1,99 +1,32 @@
""" """
Readline helper functions for dream.py (linux and mac only). 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 os
import re import re
import atexit import atexit
completer = None
# ---------------readline utilities--------------------- # ---------------readline utilities---------------------
try: try:
import readline import readline
readline_available = True readline_available = True
except: except:
readline_available = False readline_available = False
#to simulate what happens on windows systems, uncomment
# this line
#readline_available = False
class Completer: IMG_EXTENSIONS = ('.png','.jpg','.jpeg')
def __init__(self, options): COMMANDS = (
self.options = sorted(options)
return
def complete(self, text, state):
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:
# This is the first time for this text, so build a match list.
if text:
self.matches = [
s for s in self.options if s and s.startswith(text)
]
else:
self.matches = self.options[:]
# Return the state'th item from the match list,
# if we have that many.
try:
response = self.matches[state]
except IndexError:
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()
else:
path = text
matches = list()
path = os.path.expanduser(path)
if len(path) == 0:
matches.append(text + './')
else:
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))
try:
response = matches[state]
except IndexError:
response = None
return response
if readline_available:
readline.set_completer(
Completer(
[
'--steps','-s', '--steps','-s',
'--seed','-S', '--seed','-S',
'--iterations','-n', '--iterations','-n',
@ -115,12 +48,220 @@ if readline_available:
'--upscale','-U', '--upscale','-U',
'-save_orig','--save_original', '-save_orig','--save_original',
'--skip_normalize','-x', '--skip_normalize','-x',
'--log_tokenization','t', '--log_tokenization','-t',
] '!fix','!fetch',
).complete
) )
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.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 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.
elif text:
self.matches = [
s for s in self.options if s and s.startswith(text)
]
else:
self.matches = self.options[:]
# Return the state'th item from the match list,
# if we have that many.
try:
response = self.matches[state]
except IndexError:
response = None
return response
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:
switch = ''
partial = text
matches = list()
for s in self.seeds:
if s.startswith(partial):
matches.append(switch+s)
matches.sort()
return matches
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)
else:
dir = ''
path= os.path.join(dir,path)
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.complete
)
readline.set_auto_history(False)
readline.set_pre_input_hook(completer._pre_input_hook)
readline.set_completer_delims(' ') readline.set_completer_delims(' ')
readline.parse_and_bind('tab: complete') 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') histfile = os.path.join(os.path.expanduser('~'), '.dream_history')
try: try:
@ -129,3 +270,6 @@ if readline_available:
except FileNotFoundError: except FileNotFoundError:
pass pass
atexit.register(readline.write_history_file, histfile) atexit.register(readline.write_history_file, histfile)
else:
completer = DummyCompleter(COMMANDS)

26
ldm/generate.py
View File

@ -500,25 +500,26 @@ class Generate:
opt = None, opt = None,
): ):
# retrieve the seed from the image; # 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 seed = None
image_metadata = None image_metadata = None
prompt = None prompt = None
try:
args = metadata_from_png(image_path) args = metadata_from_png(image_path)
seed = args.seed seed = args.seed
prompt = args.prompt prompt = args.prompt
print(f'>> retrieved seed {seed} and prompt "{prompt}" from {image_path}') 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)
if not seed: if not seed:
print('* Could not recover seed for image. Replacing with 42. This will not affect image quality') print('* Could not recover seed for image. Replacing with 42. This will not affect image quality')
seed = 42 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 # face fixers and esrgan take an Image, but embiggen takes a path
image = Image.open(image_path) image = Image.open(image_path)
@ -540,6 +541,7 @@ class Generate:
save_original = save_original, save_original = save_original,
upscale = upscale, upscale = upscale,
image_callback = callback, image_callback = callback,
prefix = prefix,
) )
elif tool == 'embiggen': elif tool == 'embiggen':
@ -726,7 +728,9 @@ class Generate:
strength = 0.0, strength = 0.0,
codeformer_fidelity = 0.75, codeformer_fidelity = 0.75,
save_original = False, save_original = False,
image_callback = None): image_callback = None,
prefix = None,
):
for r in image_list: for r in image_list:
image, seed = r image, seed = r
@ -760,7 +764,7 @@ class Generate:
) )
if image_callback is not None: if image_callback is not None:
image_callback(image, seed, upscaled=True) image_callback(image, seed, upscaled=True, use_prefix=prefix)
else: else:
r[0] = image r[0] = image
@ -869,10 +873,6 @@ class Generate:
def _create_init_image(self, image): def _create_init_image(self, image):
image = image.convert('RGB') image = image.convert('RGB')
# print(
# f'>> DEBUG: writing the image to img.png'
# )
# image.save('img.png')
image = np.array(image).astype(np.float32) / 255.0 image = np.array(image).astype(np.float32) / 255.0
image = image[None].transpose(0, 3, 1, 2) image = image[None].transpose(0, 3, 1, 2)
image = torch.from_numpy(image) image = torch.from_numpy(image)

254
scripts/dream.py
View File

@ -9,18 +9,17 @@ import copy
import warnings import warnings
import time import time
sys.path.append('.') # corrects a weird problem on Macs sys.path.append('.') # corrects a weird problem on Macs
import ldm.dream.readline from ldm.dream.readline import completer
from ldm.dream.args import Args, metadata_dumps, metadata_from_png from ldm.dream.args import Args, metadata_dumps, metadata_from_png, dream_cmd_from_png
from ldm.dream.pngwriter import PngWriter from ldm.dream.pngwriter import PngWriter
from ldm.dream.server import DreamServer, ThreadingDreamServer
from ldm.dream.image_util import make_grid from ldm.dream.image_util import make_grid
from ldm.dream.log import write_log from ldm.dream.log import write_log
from omegaconf import OmegaConf from omegaconf import OmegaConf
from backend.invoke_ai_web_server import InvokeAIWebServer
# Placeholder to be replaced with proper class that tracks the # The output counter labels each output and is keyed to the
# outputs and associates with the prompt that generated them. # command-line history
# Just want to get the formatting look right for now. output_cntr = completer.get_current_history_length()+1
output_cntr = 0
def main(): def main():
"""Initialize command-line parsers and the diffusion model""" """Initialize command-line parsers and the diffusion model"""
@ -111,16 +110,16 @@ def main():
#set additional option #set additional option
gen.free_gpu_mem = opt.free_gpu_mem gen.free_gpu_mem = opt.free_gpu_mem
# web server loops forever
if opt.web:
invoke_ai_web_server_loop(gen, gfpgan, codeformer, esrgan)
sys.exit(0)
if not infile: if not infile:
print( print(
"\n* Initialization done! Awaiting your command (-h for help, 'q' to quit)" "\n* Initialization done! Awaiting your command (-h for help, 'q' to quit)"
) )
# web server loops forever
if opt.web:
dream_server_loop(gen, opt.host, opt.port, opt.outdir, gfpgan)
sys.exit(0)
main_loop(gen, opt, infile) main_loop(gen, opt, infile)
# TODO: main_loop() has gotten busy. Needs to be refactored. # TODO: main_loop() has gotten busy. Needs to be refactored.
@ -142,6 +141,9 @@ def main_loop(gen, opt, infile):
while not done: while not done:
operation = 'generate' # default operation, alternative is 'postprocess' operation = 'generate' # default operation, alternative is 'postprocess'
if completer:
completer.set_default_dir(opt.outdir)
try: try:
command = get_next_command(infile) command = get_next_command(infile)
except EOFError: except EOFError:
@ -159,17 +161,29 @@ def main_loop(gen, opt, infile):
done = True done = True
break break
if command.startswith( if command.startswith('!dream'): # in case a stored prompt still contains the !dream command
'!dream'
): # in case a stored prompt still contains the !dream command
command = command.replace('!dream ','',1) command = command.replace('!dream ','',1)
if command.startswith( if command.startswith('!fix'):
'!fix'
):
command = command.replace('!fix ','',1) command = command.replace('!fix ','',1)
operation = 'postprocess' 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: if opt.parse_cmd(command) is None:
continue continue
@ -219,37 +233,15 @@ def main_loop(gen, opt, infile):
opt.strength = 0.75 if opt.out_direction is None else 0.83 opt.strength = 0.75 if opt.out_direction is None else 0.83
if opt.with_variations is not None: if opt.with_variations is not None:
# shotgun parsing, woo opt.with_variations = split_variations(opt.with_variations)
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.prompt_as_dir: if opt.prompt_as_dir:
# sanitize the prompt to a valid folder name # sanitize the prompt to a valid folder name
subdir = path_filter.sub('_', opt.prompt)[:name_max].rstrip(' .') subdir = path_filter.sub('_', opt.prompt)[:name_max].rstrip(' .')
# truncate path to maximum allowed length # truncate path to maximum allowed length
# 27 is the length of '######.##########.##.png', plus two separators and a NUL # 39 is the length of '######.##########.##########-##.png', plus two separators and a NUL
subdir = subdir[:(path_max - 27 - len(os.path.abspath(opt.outdir)))] subdir = subdir[:(path_max - 39 - len(os.path.abspath(opt.outdir)))]
current_outdir = os.path.join(opt.outdir, subdir) current_outdir = os.path.join(opt.outdir, subdir)
print('Writing files to directory: "' + current_outdir + '"') print('Writing files to directory: "' + current_outdir + '"')
@ -266,37 +258,35 @@ def main_loop(gen, opt, infile):
last_results = [] last_results = []
try: try:
file_writer = PngWriter(current_outdir) file_writer = PngWriter(current_outdir)
prefix = file_writer.unique_prefix()
results = [] # list of filename, prompt pairs results = [] # list of filename, prompt pairs
grid_images = dict() # seed -> Image, only used if `opt.grid` grid_images = dict() # seed -> Image, only used if `opt.grid`
prior_variations = opt.with_variations or [] 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 # note the seed is the seed of the current image
# the first_seed is the original seed that noise is added to # the first_seed is the original seed that noise is added to
# when the -v switch is used to generate variations # when the -v switch is used to generate variations
path = None
nonlocal prior_variations nonlocal prior_variations
if use_prefix is not None:
prefix = use_prefix
else:
prefix = file_writer.unique_prefix()
path = None
if opt.grid: if opt.grid:
grid_images[seed] = image grid_images[seed] = image
else: else:
if operation == 'postprocess': postprocessed = upscaled if upscaled else operation=='postprocess'
filename = choose_postprocess_name(opt.prompt) filename, formatted_dream_prompt = prepare_image_metadata(
elif upscaled and opt.save_original: opt,
filename = f'{prefix}.{seed}.postprocessed.png' prefix,
else: seed,
filename = f'{prefix}.{seed}.png' operation,
if opt.variation_amount > 0: prior_variations,
first_seed = first_seed or seed postprocessed,
this_variation = [[seed, opt.variation_amount]] first_seed
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)
path = file_writer.save_image_and_prompt_to_png( path = file_writer.save_image_and_prompt_to_png(
image = image, image = image,
dream_prompt = formatted_dream_prompt, dream_prompt = formatted_dream_prompt,
@ -310,10 +300,15 @@ def main_loop(gen, opt, infile):
if (not upscaled) or opt.save_original: if (not upscaled) or opt.save_original:
# only append to results if we didn't overwrite an earlier output # only append to results if we didn't overwrite an earlier output
results.append([path, formatted_dream_prompt]) 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]) last_results.append([path, seed])
if operation == 'generate': if operation == 'generate':
catch_ctrl_c = infile is None # if running interactively, we catch keyboard interrupts catch_ctrl_c = infile is None # if running interactively, we catch keyboard interrupts
opt.last_operation='generate'
gen.prompt2image( gen.prompt2image(
image_callback=image_writer, image_callback=image_writer,
catch_interrupts=catch_ctrl_c, catch_interrupts=catch_ctrl_c,
@ -321,7 +316,7 @@ def main_loop(gen, opt, infile):
) )
elif operation == 'postprocess': elif operation == 'postprocess':
print(f'>> fixing {opt.prompt}') 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: if opt.grid and len(grid_images) > 0:
grid_img = make_grid(list(grid_images.values())) grid_img = make_grid(list(grid_images.values()))
@ -356,6 +351,10 @@ def main_loop(gen, opt, infile):
global output_cntr global output_cntr
output_cntr = write_log(results, log_path ,('txt', 'md'), output_cntr) output_cntr = write_log(results, log_path ,('txt', 'md'), output_cntr)
print() print()
if operation == 'postprocess':
completer.add_history(f'!fix {command}')
else:
completer.add_history(command)
print('goodbye!') print('goodbye!')
@ -377,8 +376,9 @@ def do_postprocess (gen, opt, callback):
elif opt.out_direction: elif opt.out_direction:
tool = 'outpaint' tool = 'outpaint'
opt.save_original = True # do not overwrite old image! opt.save_original = True # do not overwrite old image!
return gen.apply_postprocessor( opt.last_operation = f'postprocess:{tool}'
image_path = opt.prompt, gen.apply_postprocessor(
image_path = file_path,
tool = tool, tool = tool,
gfpgan_strength = opt.gfpgan_strength, gfpgan_strength = opt.gfpgan_strength,
codeformer_fidelity = opt.codeformer_fidelity, codeformer_fidelity = opt.codeformer_fidelity,
@ -388,18 +388,54 @@ def do_postprocess (gen, opt, callback):
callback = callback, callback = callback,
opt = opt, opt = opt,
) )
return opt.last_operation
def choose_postprocess_name(original_filename): def prepare_image_metadata(
basename,_ = os.path.splitext(os.path.basename(original_filename)) opt,
if re.search('\d+\.\d+$',basename): prefix,
return f'{basename}.fixed.png' seed,
match = re.search('(\d+\.\d+)\.fixed(-(\d+))?$',basename) operation='generate',
if match: prior_variations=[],
counter = match.group(3) or 0 postprocessed=False,
return '{prefix}-{counter:02d}.png'.format(prefix=match.group(1), counter=int(counter)+1) first_seed=None
):
if postprocessed and opt.save_original:
filename = choose_postprocess_name(opt,prefix,seed)
else: 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 def get_next_command(infile=None) -> str: # command string
if infile is None: if infile is None:
@ -414,46 +450,60 @@ def get_next_command(infile=None) -> str: # command string
print(f'#{command}') print(f'#{command}')
return command return command
def dream_server_loop(gen, host, port, outdir, gfpgan): def invoke_ai_web_server_loop(gen, gfpgan, codeformer, esrgan):
print('\n* --web was specified, starting web server...') print('\n* --web was specified, starting web server...')
# Change working directory to the stable-diffusion directory # Change working directory to the stable-diffusion directory
os.chdir( os.chdir(
os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
) )
# Start server invoke_ai_web_server = InvokeAIWebServer(generate=gen, gfpgan=gfpgan, codeformer=codeformer, esrgan=esrgan)
DreamServer.model = gen # misnomer in DreamServer - this is not the model you are looking for
DreamServer.outdir = outdir
DreamServer.gfpgan_model_exists = False
if gfpgan is not None:
DreamServer.gfpgan_model_exists = gfpgan.gfpgan_model_exists
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: try:
dream_server.serve_forever() invoke_ai_web_server.run()
except KeyboardInterrupt: except KeyboardInterrupt:
pass pass
dream_server.server_close()
def write_log_message(results, log_path): def split_variations(variations_string) -> list:
"""logs the name of the output image, prompt, and prompt args to the terminal and log file""" # shotgun parsing, woo
global output_cntr parts = []
log_lines = [f'{path}: {prompt}\n' for path, prompt in results] broken = False # python doesn't have labeled loops...
for l in log_lines: for part in variations_string.split(','):
output_cntr += 1 seed_and_weight = part.split(':')
print(f'[{output_cntr}] {l}',end='') 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: def retrieve_dream_command(opt,file_path):
file.writelines(log_lines) '''
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__': if __name__ == '__main__':
main() main()