Skip to content

Commit

Permalink
Merge branch 'development' of github.com:deNULL/stable-diffusion into…
Browse files Browse the repository at this point in the history 
… deNULL-development
  • Loading branch information
@lstein
lstein committed Sep 9, 2022
2 parents 6531446 + 171f8db commit afee7f9
Show file tree
Hide file tree
Showing 3 changed files with 100 additions and 115 deletions.
175 changes: 82 additions & 93 deletions ldm/dream/server.py
View file
@@ -1,11 +1,59 @@
import argparse
import json
import base64
import mimetypes
import os
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from ldm.dream.pngwriter import PngWriter
from ldm.dream.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', seed 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

Expand Down Expand Up @@ -81,57 +129,15 @@ def do_POST(self):

content_length = int(self.headers['Content-Length'])
post_data = json.loads(self.rfile.read(content_length))
prompt = post_data['prompt']
initimg = post_data['initimg']
strength = float(post_data['strength'])
iterations = int(post_data['iterations'])
steps = int(post_data['steps'])
width = int(post_data['width'])
height = int(post_data['height'])
fit = 'fit' in post_data
seamless = 'seamless' in post_data
cfgscale = float(post_data['cfgscale'])
sampler_name = post_data['sampler']
variation_amount = float(post_data['variation_amount'])
with_variations = post_data['with_variations']
gfpgan_strength = float(post_data['gfpgan_strength']) if gfpgan_model_exists else 0
upscale_level = post_data['upscale_level']
upscale_strength = post_data['upscale_strength']
upscale = [int(upscale_level),float(upscale_strength)] if upscale_level != '' else None
progress_images = 'progress_images' in post_data
seed = None if int(post_data['seed']) == -1 else int(post_data['seed'])

if with_variations != '':
parts = []
broken = False
for part in 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:
vseed = int(seed_and_weight[0])
vweight = float(seed_and_weight[1])
except ValueError:
print(f'could not parse with_variation part "{part}"')
broken = True
break
parts.append([vseed, vweight])
if broken:
raise CanceledException
if len(parts) > 0:
with_variations = parts
else:
with_variations = None
opt = build_opt(post_data, self.model.seed, gfpgan_model_exists)

self.canceled.clear()
print(f">> Request to generate with prompt: {prompt}")
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','')
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
Expand All @@ -144,9 +150,21 @@ def do_POST(self):
# entry should not be inserted into the image list.
def image_done(image, seed, upscaled=False):
name = f'{prefix}.{seed}.png'
path = pngwriter.save_image_and_prompt_to_png(image, f'{prompt} -S{seed}', name)

config['seed'] = seed
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:
Expand All @@ -157,24 +175,24 @@ def image_done(image, seed, upscaled=False):
) + '\n',"utf-8"))

# control state of the "postprocessing..." message
upscaling_requested = upscale or gfpgan_strength>0
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
images_generated += 1
if upscaling_requested:
action = None
if images_generated >= iterations:
if images_upscaled < iterations:
if images_generated >= opt.iterations:
if images_upscaled < opt.iterations:
action = 'upscaling-started'
else:
action = 'upscaling-done'
if action:
x = images_upscaled+1
x = images_upscaled + 1
self.wfile.write(bytes(json.dumps(
{'event':action,'processed_file_cnt':f'{x}/{iterations}'}
{'event': action, 'processed_file_cnt': f'{x}/{opt.iterations}'}
) + '\n',"utf-8"))

step_writer = PngWriter(os.path.join(self.outdir, "intermediates"))
Expand All @@ -187,60 +205,31 @@ def image_progress(sample, step):
# 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 progress_images and step % 5 == 0 and step < steps - 1:
if opt.progress_images and step % 5 == 0 and step < opt.steps - 1:
image = self.model.sample_to_image(sample)
name = f'{prefix}.{seed}.{step_index}.png'
metadata = f'{prompt} -S{seed} [intermediate]'
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 initimg is None:
if opt.init_img is None:
# Run txt2img
self.model.prompt2image(prompt,
iterations=iterations,
cfg_scale = cfgscale,
width = width,
height = height,
seed = seed,
steps = steps,
variation_amount = variation_amount,
with_variations = with_variations,
gfpgan_strength = gfpgan_strength,
upscale = upscale,
sampler_name = sampler_name,
seamless = seamless,
step_callback=image_progress,
image_callback=image_done)
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 = initimg.split(",")[1] # Ignore mime type
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(prompt,
init_img = "./img2img-tmp.png",
strength = strength,
iterations = iterations,
cfg_scale = cfgscale,
seed = seed,
steps = steps,
variation_amount = variation_amount,
with_variations = with_variations,
sampler_name = sampler_name,
width = width,
height = height,
fit = fit,
seamless = seamless,
gfpgan_strength=gfpgan_strength,
upscale = upscale,
step_callback=image_progress,
image_callback=image_done)
self.model.prompt2image(**vars(opt1), step_callback=image_progress, image_callback=image_done)
finally:
# Remove the temp file
os.remove("./img2img-tmp.png")
Expand Down
12 changes: 6 additions & 6 deletions static/dream_web/index.html
View file
Expand Up @@ -30,21 +30,21 @@ <h1>Stable Diffusion Dream Server</h1>
<input value="1" type="number" id="iterations" name="iterations" size="4">
<label for="steps">Steps:</label>
<input value="50" type="number" id="steps" name="steps">
<label for="cfgscale">Cfg Scale:</label>
<input value="7.5" type="number" id="cfgscale" name="cfgscale" step="any">
<label for="sampler">Sampler:</label>
<select id="sampler" name="sampler" value="k_lms">
<label for="cfg_scale">Cfg Scale:</label>
<input value="7.5" type="number" id="cfg_scale" name="cfg_scale" step="any">
<label for="sampler_name">Sampler:</label>
<select id="sampler_name" name="sampler_name" value="k_lms">
<option value="ddim">DDIM</option>
<option value="plms">PLMS</option>
<option value="k_lms" selected>KLMS</option>
<option value="k_dpm_2">KDPM_2</option>
<option value="k_dpm_2_a">KDPM_2A</option>
<option value="k_euler">KEULER</option>
<option value="k_euler_a">KEULER_A</option>
<option value="k_euler_a">KEULER_A</option>
<option value="k_heun">KHEUN</option>
</select>
<input type="checkbox" name="seamless" id="seamless">
<label for="seamless">Seamless circular tiling</label>
<label for="seamless">Seamless circular tiling</label>
<br>
<label title="Set to multiple of 64" for="width">Width:</label>
<select id="width" name="width" value="512">
Expand Down
28 changes: 12 additions & 16 deletions static/dream_web/index.js
View file
Expand Up @@ -9,7 +9,13 @@ function toBase64(file) {

function appendOutput(src, seed, config) {
let outputNode = document.createElement("figure");
let altText = seed.toString() + " | " + config.prompt;

let variations = config.with_variations;
if (config.variation_amount > 0) {
variations = (variations ? variations + ',' : '') + seed + ':' + config.variation_amount;
}
let baseseed = (config.with_variations || config.variation_amount > 0) ? config.seed : seed;
let altText = baseseed + ' | ' + (variations ? variations + ' | ' : '') + config.prompt;

// img needs width and height for lazy loading to work
const figureContents = `
Expand All @@ -25,7 +31,7 @@ function appendOutput(src, seed, config) {
`;

outputNode.innerHTML = figureContents;
let figcaption = outputNode.querySelector('figcaption')
let figcaption = outputNode.querySelector('figcaption');

// Reload image config
figcaption.addEventListener('click', () => {
Expand All @@ -34,21 +40,11 @@ function appendOutput(src, seed, config) {
if (k == 'initimg') { continue; }
form.querySelector(`*[name=${k}]`).value = config[k];
}
if (config.variation_amount > 0 || config.with_variations != '') {
document.querySelector("#seed").value = config.seed;
} else {
document.querySelector("#seed").value = seed;
}

if (config.variation_amount > 0) {
let oldVarAmt = document.querySelector("#variation_amount").value
let oldVariations = document.querySelector("#with_variations").value
let varSep = ''
document.querySelector("#variation_amount").value = 0;
if (document.querySelector("#with_variations").value != '') {
varSep = ","
}
document.querySelector("#with_variations").value = oldVariations + varSep + seed + ':' + config.variation_amount
document.querySelector("#seed").value = baseseed;
document.querySelector("#with_variations").value = variations || '';
if (document.querySelector("#variation_amount").value <= 0) {
document.querySelector("#variation_amount").value = 0.2;
}

saveFields(document.querySelector("#generate-form"));
Expand Down

0 comments on commit afee7f9

Please sign in to comment.