Merge branch 'master' of https://github.com/AUTOMATIC1111/stable-diffusion-webui into Base
JashoBell
3 years ago
commit
d2c7ad2fec
@ -0,0 +1,77 @@
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class MemUsageMonitor(threading.Thread):
|
||||
run_flag = None
|
||||
device = None
|
||||
disabled = False
|
||||
opts = None
|
||||
data = None
|
||||
|
||||
def __init__(self, name, device, opts):
|
||||
threading.Thread.__init__(self)
|
||||
self.name = name
|
||||
self.device = device
|
||||
self.opts = opts
|
||||
|
||||
self.daemon = True
|
||||
self.run_flag = threading.Event()
|
||||
self.data = defaultdict(int)
|
||||
|
||||
def run(self):
|
||||
if self.disabled:
|
||||
return
|
||||
|
||||
while True:
|
||||
self.run_flag.wait()
|
||||
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
self.data.clear()
|
||||
|
||||
if self.opts.memmon_poll_rate <= 0:
|
||||
self.run_flag.clear()
|
||||
continue
|
||||
|
||||
self.data["min_free"] = torch.cuda.mem_get_info()[0]
|
||||
|
||||
while self.run_flag.is_set():
|
||||
free, total = torch.cuda.mem_get_info() # calling with self.device errors, torch bug?
|
||||
self.data["min_free"] = min(self.data["min_free"], free)
|
||||
|
||||
time.sleep(1 / self.opts.memmon_poll_rate)
|
||||
|
||||
def dump_debug(self):
|
||||
print(self, 'recorded data:')
|
||||
for k, v in self.read().items():
|
||||
print(k, -(v // -(1024 ** 2)))
|
||||
|
||||
print(self, 'raw torch memory stats:')
|
||||
tm = torch.cuda.memory_stats(self.device)
|
||||
for k, v in tm.items():
|
||||
if 'bytes' not in k:
|
||||
continue
|
||||
print('\t' if 'peak' in k else '', k, -(v // -(1024 ** 2)))
|
||||
|
||||
print(torch.cuda.memory_summary())
|
||||
|
||||
def monitor(self):
|
||||
self.run_flag.set()
|
||||
|
||||
def read(self):
|
||||
free, total = torch.cuda.mem_get_info()
|
||||
self.data["total"] = total
|
||||
|
||||
torch_stats = torch.cuda.memory_stats(self.device)
|
||||
self.data["active_peak"] = torch_stats["active_bytes.all.peak"]
|
||||
self.data["reserved_peak"] = torch_stats["reserved_bytes.all.peak"]
|
||||
self.data["system_peak"] = total - self.data["min_free"]
|
||||
|
||||
return self.data
|
||||
|
||||
def stop(self):
|
||||
self.run_flag.clear()
|
||||
return self.read()
|
||||
@ -0,0 +1,148 @@
|
||||
import glob
|
||||
import os.path
|
||||
import sys
|
||||
from collections import namedtuple
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
|
||||
from ldm.util import instantiate_from_config
|
||||
|
||||
from modules import shared
|
||||
|
||||
CheckpointInfo = namedtuple("CheckpointInfo", ['filename', 'title', 'hash'])
|
||||
checkpoints_list = {}
|
||||
|
||||
try:
|
||||
# this silences the annoying "Some weights of the model checkpoint were not used when initializing..." message at start.
|
||||
|
||||
from transformers import logging
|
||||
|
||||
logging.set_verbosity_error()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def list_models():
|
||||
checkpoints_list.clear()
|
||||
|
||||
model_dir = os.path.abspath(shared.cmd_opts.ckpt_dir)
|
||||
|
||||
def modeltitle(path, h):
|
||||
abspath = os.path.abspath(path)
|
||||
|
||||
if abspath.startswith(model_dir):
|
||||
name = abspath.replace(model_dir, '')
|
||||
else:
|
||||
name = os.path.basename(path)
|
||||
|
||||
if name.startswith("\\") or name.startswith("/"):
|
||||
name = name[1:]
|
||||
|
||||
return f'{name} [{h}]'
|
||||
|
||||
cmd_ckpt = shared.cmd_opts.ckpt
|
||||
if os.path.exists(cmd_ckpt):
|
||||
h = model_hash(cmd_ckpt)
|
||||
title = modeltitle(cmd_ckpt, h)
|
||||
checkpoints_list[title] = CheckpointInfo(cmd_ckpt, title, h)
|
||||
elif cmd_ckpt is not None and cmd_ckpt != shared.default_sd_model_file:
|
||||
print(f"Checkpoint in --ckpt argument not found: {cmd_ckpt}", file=sys.stderr)
|
||||
|
||||
if os.path.exists(model_dir):
|
||||
for filename in glob.glob(model_dir + '/**/*.ckpt', recursive=True):
|
||||
h = model_hash(filename)
|
||||
title = modeltitle(filename, h)
|
||||
checkpoints_list[title] = CheckpointInfo(filename, title, h)
|
||||
|
||||
|
||||
def model_hash(filename):
|
||||
try:
|
||||
with open(filename, "rb") as file:
|
||||
import hashlib
|
||||
m = hashlib.sha256()
|
||||
|
||||
file.seek(0x100000)
|
||||
m.update(file.read(0x10000))
|
||||
return m.hexdigest()[0:8]
|
||||
except FileNotFoundError:
|
||||
return 'NOFILE'
|
||||
|
||||
|
||||
def select_checkpoint():
|
||||
model_checkpoint = shared.opts.sd_model_checkpoint
|
||||
checkpoint_info = checkpoints_list.get(model_checkpoint, None)
|
||||
if checkpoint_info is not None:
|
||||
return checkpoint_info
|
||||
|
||||
if len(checkpoints_list) == 0:
|
||||
print(f"Checkpoint {model_checkpoint} not found and no other checkpoints found", file=sys.stderr)
|
||||
return None
|
||||
|
||||
checkpoint_info = next(iter(checkpoints_list.values()))
|
||||
if model_checkpoint is not None:
|
||||
print(f"Checkpoint {model_checkpoint} not found; loading fallback {checkpoint_info.title}", file=sys.stderr)
|
||||
|
||||
return checkpoint_info
|
||||
|
||||
|
||||
def load_model_weights(model, checkpoint_file, sd_model_hash):
|
||||
print(f"Loading weights [{sd_model_hash}] from {checkpoint_file}")
|
||||
|
||||
pl_sd = torch.load(checkpoint_file, map_location="cpu")
|
||||
if "global_step" in pl_sd:
|
||||
print(f"Global Step: {pl_sd['global_step']}")
|
||||
sd = pl_sd["state_dict"]
|
||||
|
||||
model.load_state_dict(sd, strict=False)
|
||||
|
||||
if shared.cmd_opts.opt_channelslast:
|
||||
model.to(memory_format=torch.channels_last)
|
||||
|
||||
if not shared.cmd_opts.no_half:
|
||||
model.half()
|
||||
|
||||
model.sd_model_hash = sd_model_hash
|
||||
model.sd_model_checkpint = checkpoint_file
|
||||
|
||||
|
||||
def load_model():
|
||||
from modules import lowvram, sd_hijack
|
||||
checkpoint_info = select_checkpoint()
|
||||
|
||||
sd_config = OmegaConf.load(shared.cmd_opts.config)
|
||||
sd_model = instantiate_from_config(sd_config.model)
|
||||
load_model_weights(sd_model, checkpoint_info.filename, checkpoint_info.hash)
|
||||
|
||||
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
|
||||
lowvram.setup_for_low_vram(sd_model, shared.cmd_opts.medvram)
|
||||
else:
|
||||
sd_model.to(shared.device)
|
||||
|
||||
sd_hijack.model_hijack.hijack(sd_model)
|
||||
|
||||
sd_model.eval()
|
||||
|
||||
print(f"Model loaded.")
|
||||
return sd_model
|
||||
|
||||
|
||||
def reload_model_weights(sd_model, info=None):
|
||||
from modules import lowvram, devices
|
||||
checkpoint_info = info or select_checkpoint()
|
||||
|
||||
if sd_model.sd_model_checkpint == checkpoint_info.filename:
|
||||
return
|
||||
|
||||
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
|
||||
lowvram.send_everything_to_cpu()
|
||||
else:
|
||||
sd_model.to(devices.cpu)
|
||||
|
||||
load_model_weights(sd_model, checkpoint_info.filename, checkpoint_info.hash)
|
||||
|
||||
if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:
|
||||
sd_model.to(devices.device)
|
||||
|
||||
print(f"Weights loaded.")
|
||||
return sd_model
|
||||
@ -0,0 +1,290 @@
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import skimage
|
||||
|
||||
import modules.scripts as scripts
|
||||
import gradio as gr
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
from modules import images, processing, devices
|
||||
from modules.processing import Processed, process_images
|
||||
from modules.shared import opts, cmd_opts, state
|
||||
|
||||
|
||||
def expand(x, dir, amount, power=0.75):
|
||||
is_left = dir == 3
|
||||
is_right = dir == 1
|
||||
is_up = dir == 0
|
||||
is_down = dir == 2
|
||||
|
||||
if is_left or is_right:
|
||||
noise = np.zeros((x.shape[0], amount, 3), dtype=float)
|
||||
indexes = np.random.random((x.shape[0], amount)) ** power * (1 - np.arange(amount) / amount)
|
||||
if is_right:
|
||||
indexes = 1 - indexes
|
||||
indexes = (indexes * (x.shape[1] - 1)).astype(int)
|
||||
|
||||
for row in range(x.shape[0]):
|
||||
if is_left:
|
||||
noise[row] = x[row][indexes[row]]
|
||||
else:
|
||||
noise[row] = np.flip(x[row][indexes[row]], axis=0)
|
||||
|
||||
x = np.concatenate([noise, x] if is_left else [x, noise], axis=1)
|
||||
return x
|
||||
|
||||
if is_up or is_down:
|
||||
noise = np.zeros((amount, x.shape[1], 3), dtype=float)
|
||||
indexes = np.random.random((x.shape[1], amount)) ** power * (1 - np.arange(amount) / amount)
|
||||
if is_down:
|
||||
indexes = 1 - indexes
|
||||
indexes = (indexes * x.shape[0] - 1).astype(int)
|
||||
|
||||
for row in range(x.shape[1]):
|
||||
if is_up:
|
||||
noise[:, row] = x[:, row][indexes[row]]
|
||||
else:
|
||||
noise[:, row] = np.flip(x[:, row][indexes[row]], axis=0)
|
||||
|
||||
x = np.concatenate([noise, x] if is_up else [x, noise], axis=0)
|
||||
return x
|
||||
|
||||
|
||||
def get_matched_noise(_np_src_image, np_mask_rgb, noise_q=1, color_variation=0.05):
|
||||
# helper fft routines that keep ortho normalization and auto-shift before and after fft
|
||||
def _fft2(data):
|
||||
if data.ndim > 2: # has channels
|
||||
out_fft = np.zeros((data.shape[0], data.shape[1], data.shape[2]), dtype=np.complex128)
|
||||
for c in range(data.shape[2]):
|
||||
c_data = data[:, :, c]
|
||||
out_fft[:, :, c] = np.fft.fft2(np.fft.fftshift(c_data), norm="ortho")
|
||||
out_fft[:, :, c] = np.fft.ifftshift(out_fft[:, :, c])
|
||||
else: # one channel
|
||||
out_fft = np.zeros((data.shape[0], data.shape[1]), dtype=np.complex128)
|
||||
out_fft[:, :] = np.fft.fft2(np.fft.fftshift(data), norm="ortho")
|
||||
out_fft[:, :] = np.fft.ifftshift(out_fft[:, :])
|
||||
|
||||
return out_fft
|
||||
|
||||
def _ifft2(data):
|
||||
if data.ndim > 2: # has channels
|
||||
out_ifft = np.zeros((data.shape[0], data.shape[1], data.shape[2]), dtype=np.complex128)
|
||||
for c in range(data.shape[2]):
|
||||
c_data = data[:, :, c]
|
||||
out_ifft[:, :, c] = np.fft.ifft2(np.fft.fftshift(c_data), norm="ortho")
|
||||
out_ifft[:, :, c] = np.fft.ifftshift(out_ifft[:, :, c])
|
||||
else: # one channel
|
||||
out_ifft = np.zeros((data.shape[0], data.shape[1]), dtype=np.complex128)
|
||||
out_ifft[:, :] = np.fft.ifft2(np.fft.fftshift(data), norm="ortho")
|
||||
out_ifft[:, :] = np.fft.ifftshift(out_ifft[:, :])
|
||||
|
||||
return out_ifft
|
||||
|
||||
def _get_gaussian_window(width, height, std=3.14, mode=0):
|
||||
window_scale_x = float(width / min(width, height))
|
||||
window_scale_y = float(height / min(width, height))
|
||||
|
||||
window = np.zeros((width, height))
|
||||
x = (np.arange(width) / width * 2. - 1.) * window_scale_x
|
||||
for y in range(height):
|
||||
fy = (y / height * 2. - 1.) * window_scale_y
|
||||
if mode == 0:
|
||||
window[:, y] = np.exp(-(x ** 2 + fy ** 2) * std)
|
||||
else:
|
||||
window[:, y] = (1 / ((x ** 2 + 1.) * (fy ** 2 + 1.))) ** (std / 3.14) # hey wait a minute that's not gaussian
|
||||
|
||||
return window
|
||||
|
||||
def _get_masked_window_rgb(np_mask_grey, hardness=1.):
|
||||
np_mask_rgb = np.zeros((np_mask_grey.shape[0], np_mask_grey.shape[1], 3))
|
||||
if hardness != 1.:
|
||||
hardened = np_mask_grey[:] ** hardness
|
||||
else:
|
||||
hardened = np_mask_grey[:]
|
||||
for c in range(3):
|
||||
np_mask_rgb[:, :, c] = hardened[:]
|
||||
return np_mask_rgb
|
||||
|
||||
width = _np_src_image.shape[0]
|
||||
height = _np_src_image.shape[1]
|
||||
num_channels = _np_src_image.shape[2]
|
||||
|
||||
np_src_image = _np_src_image[:] * (1. - np_mask_rgb)
|
||||
np_mask_grey = (np.sum(np_mask_rgb, axis=2) / 3.)
|
||||
img_mask = np_mask_grey > 1e-6
|
||||
ref_mask = np_mask_grey < 1e-3
|
||||
|
||||
windowed_image = _np_src_image * (1. - _get_masked_window_rgb(np_mask_grey))
|
||||
windowed_image /= np.max(windowed_image)
|
||||
windowed_image += np.average(_np_src_image) * np_mask_rgb # / (1.-np.average(np_mask_rgb)) # rather than leave the masked area black, we get better results from fft by filling the average unmasked color
|
||||
|
||||
src_fft = _fft2(windowed_image) # get feature statistics from masked src img
|
||||
src_dist = np.absolute(src_fft)
|
||||
src_phase = src_fft / src_dist
|
||||
|
||||
noise_window = _get_gaussian_window(width, height, mode=1) # start with simple gaussian noise
|
||||
noise_rgb = np.random.random_sample((width, height, num_channels))
|
||||
noise_grey = (np.sum(noise_rgb, axis=2) / 3.)
|
||||
noise_rgb *= color_variation # the colorfulness of the starting noise is blended to greyscale with a parameter
|
||||
for c in range(num_channels):
|
||||
noise_rgb[:, :, c] += (1. - color_variation) * noise_grey
|
||||
|
||||
noise_fft = _fft2(noise_rgb)
|
||||
for c in range(num_channels):
|
||||
noise_fft[:, :, c] *= noise_window
|
||||
noise_rgb = np.real(_ifft2(noise_fft))
|
||||
shaped_noise_fft = _fft2(noise_rgb)
|
||||
shaped_noise_fft[:, :, :] = np.absolute(shaped_noise_fft[:, :, :]) ** 2 * (src_dist ** noise_q) * src_phase # perform the actual shaping
|
||||
|
||||
brightness_variation = 0. # color_variation # todo: temporarily tieing brightness variation to color variation for now
|
||||
contrast_adjusted_np_src = _np_src_image[:] * (brightness_variation + 1.) - brightness_variation * 2.
|
||||
|
||||
# scikit-image is used for histogram matching, very convenient!
|
||||
shaped_noise = np.real(_ifft2(shaped_noise_fft))
|
||||
shaped_noise -= np.min(shaped_noise)
|
||||
shaped_noise /= np.max(shaped_noise)
|
||||
shaped_noise[img_mask, :] = skimage.exposure.match_histograms(shaped_noise[img_mask, :] ** 1., contrast_adjusted_np_src[ref_mask, :], channel_axis=1)
|
||||
shaped_noise = _np_src_image[:] * (1. - np_mask_rgb) + shaped_noise * np_mask_rgb
|
||||
|
||||
matched_noise = shaped_noise[:]
|
||||
|
||||
return np.clip(matched_noise, 0., 1.)
|
||||
|
||||
|
||||
|
||||
class Script(scripts.Script):
|
||||
def title(self):
|
||||
return "Outpainting mk2"
|
||||
|
||||
def show(self, is_img2img):
|
||||
return is_img2img
|
||||
|
||||
def ui(self, is_img2img):
|
||||
if not is_img2img:
|
||||
return None
|
||||
|
||||
info = gr.HTML("<p style=\"margin-bottom:0.75em\">Recommended settings: Sampling Steps: 80-100, Sampler: Euler a, Denoising strength: 0.8</p>")
|
||||
|
||||
pixels = gr.Slider(label="Pixels to expand", minimum=8, maximum=256, step=8, value=128)
|
||||
mask_blur = gr.Slider(label='Mask blur', minimum=0, maximum=64, step=1, value=8, visible=False)
|
||||
direction = gr.CheckboxGroup(label="Outpainting direction", choices=['left', 'right', 'up', 'down'], value=['left', 'right', 'up', 'down'])
|
||||
noise_q = gr.Slider(label="Fall-off exponent (lower=higher detail)", minimum=0.0, maximum=4.0, step=0.01, value=1.0)
|
||||
color_variation = gr.Slider(label="Color variation", minimum=0.0, maximum=1.0, step=0.01, value=0.05)
|
||||
|
||||
return [info, pixels, mask_blur, direction, noise_q, color_variation]
|
||||
|
||||
def run(self, p, _, pixels, mask_blur, direction, noise_q, color_variation):
|
||||
initial_seed_and_info = [None, None]
|
||||
|
||||
process_width = p.width
|
||||
process_height = p.height
|
||||
|
||||
p.mask_blur = mask_blur*4
|
||||
p.inpaint_full_res = False
|
||||
p.inpainting_fill = 1
|
||||
p.do_not_save_samples = True
|
||||
p.do_not_save_grid = True
|
||||
|
||||
left = pixels if "left" in direction else 0
|
||||
right = pixels if "right" in direction else 0
|
||||
up = pixels if "up" in direction else 0
|
||||
down = pixels if "down" in direction else 0
|
||||
|
||||
init_img = p.init_images[0]
|
||||
target_w = math.ceil((init_img.width + left + right) / 64) * 64
|
||||
target_h = math.ceil((init_img.height + up + down) / 64) * 64
|
||||
|
||||
if left > 0:
|
||||
left = left * (target_w - init_img.width) // (left + right)
|
||||
if right > 0:
|
||||
right = target_w - init_img.width - left
|
||||
|
||||
if up > 0:
|
||||
up = up * (target_h - init_img.height) // (up + down)
|
||||
|
||||
if down > 0:
|
||||
down = target_h - init_img.height - up
|
||||
|
||||
init_image = p.init_images[0]
|
||||
|
||||
state.job_count = (1 if left > 0 else 0) + (1 if right > 0 else 0)+ (1 if up > 0 else 0)+ (1 if down > 0 else 0)
|
||||
|
||||
def expand(init, expand_pixels, is_left=False, is_right=False, is_top=False, is_bottom=False):
|
||||
is_horiz = is_left or is_right
|
||||
is_vert = is_top or is_bottom
|
||||
pixels_horiz = expand_pixels if is_horiz else 0
|
||||
pixels_vert = expand_pixels if is_vert else 0
|
||||
|
||||
img = Image.new("RGB", (init.width + pixels_horiz, init.height + pixels_vert))
|
||||
img.paste(init, (pixels_horiz if is_left else 0, pixels_vert if is_top else 0))
|
||||
mask = Image.new("RGB", (init.width + pixels_horiz, init.height + pixels_vert), "white")
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle((
|
||||
expand_pixels + mask_blur if is_left else 0,
|
||||
expand_pixels + mask_blur if is_top else 0,
|
||||
mask.width - expand_pixels - mask_blur if is_right else mask.width,
|
||||
mask.height - expand_pixels - mask_blur if is_bottom else mask.height,
|
||||
), fill="black")
|
||||
|
||||
np_image = (np.asarray(img) / 255.0).astype(np.float64)
|
||||
np_mask = (np.asarray(mask) / 255.0).astype(np.float64)
|
||||
noised = get_matched_noise(np_image, np_mask, noise_q, color_variation)
|
||||
out = Image.fromarray(np.clip(noised * 255., 0., 255.).astype(np.uint8), mode="RGB")
|
||||
|
||||
target_width = min(process_width, init.width + pixels_horiz) if is_horiz else img.width
|
||||
target_height = min(process_height, init.height + pixels_vert) if is_vert else img.height
|
||||
|
||||
crop_region = (
|
||||
0 if is_left else out.width - target_width,
|
||||
0 if is_top else out.height - target_height,
|
||||
target_width if is_left else out.width,
|
||||
target_height if is_top else out.height,
|
||||
)
|
||||
|
||||
image_to_process = out.crop(crop_region)
|
||||
mask = mask.crop(crop_region)
|
||||
|
||||
p.width = target_width if is_horiz else img.width
|
||||
p.height = target_height if is_vert else img.height
|
||||
p.init_images = [image_to_process]
|
||||
p.image_mask = mask
|
||||
|
||||
latent_mask = Image.new("RGB", (p.width, p.height), "white")
|
||||
draw = ImageDraw.Draw(latent_mask)
|
||||
draw.rectangle((
|
||||
expand_pixels + mask_blur * 2 if is_left else 0,
|
||||
expand_pixels + mask_blur * 2 if is_top else 0,
|
||||
mask.width - expand_pixels - mask_blur * 2 if is_right else mask.width,
|
||||
mask.height - expand_pixels - mask_blur * 2 if is_bottom else mask.height,
|
||||
), fill="black")
|
||||
p.latent_mask = latent_mask
|
||||
|
||||
proc = process_images(p)
|
||||
proc_img = proc.images[0]
|
||||
|
||||
if initial_seed_and_info[0] is None:
|
||||
initial_seed_and_info[0] = proc.seed
|
||||
initial_seed_and_info[1] = proc.info
|
||||
|
||||
out.paste(proc_img, (0 if is_left else out.width - proc_img.width, 0 if is_top else out.height - proc_img.height))
|
||||
return out
|
||||
|
||||
img = init_image
|
||||
|
||||
if left > 0:
|
||||
img = expand(img, left, is_left=True)
|
||||
if right > 0:
|
||||
img = expand(img, right, is_right=True)
|
||||
if up > 0:
|
||||
img = expand(img, up, is_top=True)
|
||||
if down > 0:
|
||||
img = expand(img, down, is_bottom=True)
|
||||
|
||||
res = Processed(p, [img], initial_seed_and_info[0], initial_seed_and_info[1])
|
||||
|
||||
if opts.samples_save:
|
||||
images.save_image(img, p.outpath_samples, "", res.seed, p.prompt, opts.grid_format, info=res.info, p=p)
|
||||
|
||||
return res
|
||||
|
||||
Loading…
Reference in New Issue