From 61aa82dcd00d41b05de3affb5ea0d76bc52e1167 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 6 Mar 2023 16:52:48 -0500 Subject: [PATCH] add sdp cross-optimization --- automatic.sh | 1 + cli/modules/process.py | 99 ++++++++++++++++++++---------- cli/modules/sdapi.py | 15 ++++- cli/train-lora.py | 37 +++++++---- modules/sd_hijack.py | 4 ++ modules/sd_hijack_optimizations.py | 42 +++++++++++++ modules/shared.py | 1 + 7 files changed, 153 insertions(+), 46 deletions(-) diff --git a/automatic.sh b/automatic.sh index da1151dc6..0b0841672 100755 --- a/automatic.sh +++ b/automatic.sh @@ -14,6 +14,7 @@ if [ "$PYTHON" == "" ]; then fi CMD="launch.py --api --xformers --disable-console-progressbars --gradio-queue --skip-version-check --skip-install --skip-torch-cuda-test --disable-nan-check --theme dark --cors-allow-origins=http://127.0.0.1:7860" +# CMD="launch.py --api --opt-sdp-attention --disable-console-progressbars --gradio-queue --skip-version-check --skip-install --skip-torch-cuda-test --disable-nan-check --theme dark --cors-allow-origins=http://127.0.0.1:7860" MODE=optimized diff --git a/cli/modules/process.py b/cli/modules/process.py index 28b9f699c..85709bc83 100755 --- a/cli/modules/process.py +++ b/cli/modules/process.py @@ -53,25 +53,29 @@ params = Map({ 'blur_samplesize': 60, # sample size to use for blur detection 'similarity_size': 64, # base similarity detection on reduced images # original image processing settings - 'keep_original': True, # keep original image + 'keep_original': False, # keep original image # face processing settings 'extract_face': False, # extract face from image 'face_score': 0.7, # min face detection score - 'face_pad': 0.2, # pad face image percentage + 'face_pad': 0.1, # pad face image percentage 'face_model': 1, # which face model to use 0/close-up 1/standard + 'face_blur': False, # check for body blur 'face_blur_score': 1.5, # max score for face blur detection + 'face_range': False, # check for body blur 'face_range_score': 0.15, # min score for face dynamic range detection - 'face_restore': True, # attempt to restore face quality - 'face_upscale': True, # attempt to scale small faces + 'face_restore': False, # attempt to restore face quality + 'face_upscale': False, # attempt to scale small faces 'face_segmentation': False, # segmentation enabled # body processing settings - 'extract_body': False, # extract face from image + 'extract_body': False, # extract body from image 'body_score': 0.9, # min body detection score 'body_visibility': 0.5, # min visibility score for each detected body part 'body_parts': 15, # min number of detected body parts with sufficient visibility 'body_pad': 0.2, # pad body image percentage 'body_model': 2, # body model to use 0/low 1/medium 2/high + 'body_blur': False, # check for body blur 'body_blur_score': 1.8, # max score for body blur detection + 'face_range': False, # check for body blur 'body_range_score': 0.15, # min score for body dynamic range detection 'body_segmentation': False, # segmentation enabled # similarity detection settings @@ -212,19 +216,21 @@ def extract_face(img): else: squared = cropped - blur = detect_blur(squared) - if blur > params.face_blur_score: - log.info({ 'process face skip': 'blur check fail', 'blur': blur }) - return None, True - else: - log.debug({ 'process face blur': blur }) + if params.face_blur: + blur = detect_blur(squared) + if blur > params.face_blur_score: + log.info({ 'process face skip': 'blur check fail', 'blur': blur }) + return None, True + else: + log.debug({ 'process face blur': blur }) - range = detect_dynamicrange(squared) - if range < params.face_range_score: - log.info({ 'process face skip': 'dynamic range check fail', 'range': range }) - return None, True - else: - log.debug({ 'process face dynamic range': range }) + if params.face_range: + range = detect_dynamicrange(squared) + if range < params.face_range_score: + log.info({ 'process face skip': 'dynamic range check fail', 'range': range }) + return None, True + else: + log.debug({ 'process face dynamic range': range }) similarity = detect_simmilar(squared) if similarity > params.similarity_score: @@ -275,19 +281,21 @@ def extract_body(img): else: squared = cropped - blur = detect_blur(squared) - if blur > params.body_blur_score: - log.info({ 'process body skip': 'blur check fail', 'blur': blur }) - return None, True - else: - log.debug({ 'process body blur': blur }) + if params.body_blur: + blur = detect_blur(squared) + if blur > params.body_blur_score: + log.info({ 'process body skip': 'blur check fail', 'blur': blur }) + return None, True + else: + log.debug({ 'process body blur': blur }) - range = detect_dynamicrange(squared) - if range < params.body_range_score: - log.info({ 'process body skip': 'dynamic range check fail', 'range': range }) - return None, True - else: - log.debug({ 'process body dynamic range': range }) + if params.body_range: + range = detect_dynamicrange(squared) + if range < params.body_range_score: + log.info({ 'process body skip': 'dynamic range check fail', 'range': range }) + return None, True + else: + log.debug({ 'process body dynamic range': range }) similarity = detect_simmilar(squared) if similarity > params.similarity_score: @@ -318,7 +326,7 @@ def encode(img): return encoded -def interrogate(img, fn): +def interrogate(img, fn, intag = None): if len(params.interrogate_model) == 0: return caption = '' @@ -329,16 +337,22 @@ def interrogate(img, fn): if model == 'clip': caption = res.caption if 'caption' in res else '' caption = caption.split(',')[0].replace('a ', '') + if intag is not None: + caption = intag + ', ' + caption if model == 'deepdanbooru': tag = res.caption if 'caption' in res else '' tags = tag.split(',') tags = [t.replace('(', '').replace(')', '').split(':')[0].strip() for t in tags] + if intag is not None: + for t in intag.split(',')[::-1]: + tags.insert(0, t.strip()) if params.interrogate_captions: file = fn.replace(params.format, '.txt') f = open(file, 'w') f.write(caption) f.close() - tags.insert(0, caption.split(' ')[0]) + pos = 0 if len(tags) == 0 else 1 + tags.insert(pos, caption.split(' ')[1]) if len(tags) > params.tag_limit: tags = tags[:params.tag_limit] log.info({ 'interrogate': caption, 'tags': tags }) @@ -348,7 +362,8 @@ def interrogate(img, fn): i = {} metadata = Map({}) -def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = False, txt = None): +# entry point when used as module +def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = False, txt = None, tag = None, opts = []): def save(img, f, what): i[what] = i.get(what, 0) + 1 if dst is None: @@ -365,10 +380,28 @@ def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = if not preview: img.save(os.path.join(dir, fn)) if not offline: - caption, tags = interrogate(img, os.path.join(dir, fn)) + caption, tags = interrogate(img, os.path.join(dir, fn), tag) metadata[os.path.join(parent, basename)] = { 'caption': caption, 'tags': ','.join(tags) } return fn + # overrides + if 'original' in opts: + params.keep_original = True + if 'face' in opts: + params.extract_face = True + if 'body' in opts: + params.extract_body = True + if 'blur' in opts: + params.face_blur = True + params.body_blur = True + if 'range' in opts: + params.face_range = True + params.body_range = True + if 'upscale' in opts: + params.face_upscale = True + if 'restore' in opts: + params.face_restore = True + log.info({ 'processing': f }) try: image = Image.open(f) diff --git a/cli/modules/sdapi.py b/cli/modules/sdapi.py index b03373637..e677e8dea 100755 --- a/cli/modules/sdapi.py +++ b/cli/modules/sdapi.py @@ -4,11 +4,12 @@ helper methods that creates HTTP session with managed connection pool provides async HTTP get/post methods and several helper methods """ +import sys +import json import aiohttp import asyncio import logging import requests -import sys from util import Map, log @@ -128,6 +129,12 @@ async def progress(): return res +def options(): + options = getsync('/sdapi/v1/options') + flags = getsync('/sdapi/v1/cmd-flags') + return { 'options': options, 'flags': flags } + + def shutdown(): try: postsync('/sdapi/v1/shutdown') @@ -168,6 +175,12 @@ if __name__ == "__main__": asyncio.run(interrupt()) if 'progress' in sys.argv: asyncio.run(progress()) + if 'options' in sys.argv: + opt = options() + log.debug({ 'options' }) + print(json.dumps(opt['options'], indent = 2)) + log.debug({ 'cmd-flags' }) + print(json.dumps(opt['flags'], indent = 2)) if 'shutdown' in sys.argv: shutdown() asyncio.run(close()) diff --git a/cli/train-lora.py b/cli/train-lora.py index 5df9e2c28..ed2be10f3 100755 --- a/cli/train-lora.py +++ b/cli/train-lora.py @@ -14,6 +14,7 @@ Disabled/broken: """ import os +import re import gc import sys import json @@ -127,12 +128,13 @@ def mem_stats(): if __name__ == '__main__': parser = argparse.ArgumentParser(description = 'train lora') - parser.add_argument('--model', type=str, default=None, required=True, help='original model to use a base for training') - parser.add_argument('--input', type=str, default=None, required=True, help='input folder with training images') - parser.add_argument('--output', type=str, default=None, required=True, help='lora name') + parser.add_argument('--model', type=str, default=None, required=False, help='original model to use a base for training, default: active model') + parser.add_argument('--input', '--dataset', type=str, default=None, required=True, help='input folder with training images') + parser.add_argument('--output', '--lora', type=str, default=None, required=True, help='lora name') parser.add_argument('--tag', type=str, default=None, required=False, help='primary tag') - parser.add_argument('--dir', type=str, default=None, required=True, help='folder containing lora checkpoints') + parser.add_argument('--dir', type=str, default=None, required=False, help='folder containing lora checkpoints') parser.add_argument('--interim', type=int, default=0, help = 'save interim checkpoints after n epoch') + parser.add_argument('--process', type=str, default='original', required=True, help='list of processing steps: original,face,body,blur,range,upscale,restore') parser.add_argument('--noprocess', default = False, action='store_true', help = 'skip processing and use existing input data') parser.add_argument('--notrain', default = False, action='store_true', help = 'just run processing and skip training') parser.add_argument('--nocaptions', default = False, action='store_true', help = 'skip creating captions and tags') @@ -152,19 +154,30 @@ if __name__ == '__main__': parser.add_argument('--locon', default=False, action='store_true', help = "use locon style training") parser.add_argument('--debug', default=False, action='store_true', help = "enable debug logging") args = parser.parse_args() + defaults = Map({ 'options': {}, 'flags': {} }) if args.offline else Map(modules.sdapi.options()) + if args.debug: log.setLevel(logging.DEBUG) log.debug({ 'debug': True }) + if args.model is None: + args.model = defaults.options.get('sd_model_checkpoint', None) + args.model = args.model.split(' [')[0] if args.model is not None else None + if args.dir is None: + args.dir = defaults.flags.get('lora_dir', None) + if not os.path.isabs(args.model) and args.dir is not None and not os.path.exists(args.model): + args.model = os.path.abspath(os.path.join(args.dir, os.pardir, 'Stable-diffusion', args.model)) + if args.dir is None: + args.dir = os.path.join(args.input, 'lora') if not os.path.exists(args.model) or not os.path.isfile(args.model): log.error({ 'lora cannot find model': args.model }) exit(1) - options.pretrained_model_name_or_path = args.model if not os.path.exists(args.input) or not os.path.isdir(args.input): log.error({ 'lora cannot find training dir': args.input }) exit(1) if not os.path.exists(args.dir) or not os.path.isdir(args.dir): log.error({ 'lora cannot find training dir': args.dir }) exit(1) + options.pretrained_model_name_or_path = args.model options.output_dir = args.dir options.output_name = args.output options.max_train_steps = args.steps @@ -200,9 +213,15 @@ if __name__ == '__main__': if not args.noprocess: # preprocess + processing_options = args.process.split(',') + processing_options = [opt.strip() for opt in re.split(',| ', args.process)] + log.info({ 'processing options': processing_options }) + + if os.path.exists(json_file): + os.remove(json_file) for f in files: try: - res, metadata = modules.process.process_file(f = f, dst = dir, preview = False, offline = args.offline, txt = args.dreambooth) + res, metadata = modules.process.process_file(f = f, dst = dir, preview = False, offline = args.offline, txt = args.dreambooth, tag = args.tag, opts = processing_options) if not args.dreambooth: with open(json_file, "w") as outfile: outfile.write(json.dumps(metadata, indent=2)) @@ -210,12 +229,6 @@ if __name__ == '__main__': exit(1) modules.process.unload_models() mem_stats() - if args.tag is not None: - for name, item in metadata.items(): - item['caption'] = args.tag + ',' + item['caption'] - item['tags'] = args.tag + ',' + item['tags'] - with open(json_file, "w") as outfile: - outfile.write(json.dumps(metadata, indent=2)) if not args.nolatents and not args.dreambooth: # create latents diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 794767831..76cb91209 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -42,6 +42,10 @@ def apply_optimizations(): ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.xformers_attention_forward ldm.modules.diffusionmodules.model.AttnBlock.forward = sd_hijack_optimizations.xformers_attnblock_forward optimization_method = 'xformers' + elif cmd_opts.opt_sdp_attention and (hasattr(torch.nn.functional, "scaled_dot_product_attention") and callable(getattr(torch.nn.functional, "scaled_dot_product_attention"))): + print("Applying scaled dot product cross attention optimization.") + ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.scaled_dot_product_attention_forward + optimization_method = 'sdp' elif cmd_opts.opt_sub_quad_attention: print("Applying sub-quadratic cross attention optimization.") ldm.modules.attention.CrossAttention.forward = sd_hijack_optimizations.sub_quad_attention_forward diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index c02d954c7..a324a5927 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -346,6 +346,48 @@ def xformers_attention_forward(self, x, context=None, mask=None): out = rearrange(out, 'b n h d -> b n (h d)', h=h) return self.to_out(out) +# Based on Diffusers usage of scaled dot product attention from https://github.com/huggingface/diffusers/blob/c7da8fd23359a22d0df2741688b5b4f33c26df21/src/diffusers/models/cross_attention.py +# The scaled_dot_product_attention_forward function contains parts of code under Apache-2.0 license listed under Scaled Dot Product Attention in the Licenses section of the web UI interface +def scaled_dot_product_attention_forward(self, x, context=None, mask=None): + batch_size, sequence_length, inner_dim = x.shape + + if mask is not None: + mask = self.prepare_attention_mask(mask, sequence_length, batch_size) + mask = mask.view(batch_size, self.heads, -1, mask.shape[-1]) + + h = self.heads + q_in = self.to_q(x) + context = default(context, x) + + context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context) + k_in = self.to_k(context_k) + v_in = self.to_v(context_v) + + head_dim = inner_dim // h + q = q_in.view(batch_size, -1, h, head_dim).transpose(1, 2) + k = k_in.view(batch_size, -1, h, head_dim).transpose(1, 2) + v = v_in.view(batch_size, -1, h, head_dim).transpose(1, 2) + + del q_in, k_in, v_in + + dtype = q.dtype + if shared.opts.upcast_attn: + q, k = q.float(), k.float() + + # the output of sdp = (batch, num_heads, seq_len, head_dim) + hidden_states = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False + ) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, h * head_dim) + hidden_states = hidden_states.to(dtype) + + # linear proj + hidden_states = self.to_out[0](hidden_states) + # dropout + hidden_states = self.to_out[1](hidden_states) + return hidden_states + def cross_attention_attnblock_forward(self, x): h_ = x h_ = self.norm(h_) diff --git a/modules/shared.py b/modules/shared.py index 21102d4b6..02319903a 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -107,6 +107,7 @@ parser.add_argument("--tls-certfile", type=str, help="Partially enables TLS, req parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None) parser.add_argument("--gradio-queue", action='store_true', help="Uses gradio queue; experimental option; breaks restart UI button") parser.add_argument("--skip-version-check", action='store_true', help="Do not check versions of torch and xformers") +parser.add_argument("--opt-sdp-attention", action='store_true', help="enable scaled dot product cross-attention layer optimization; requires PyTorch 2.*") parser.add_argument("--no-hashing", action='store_true', help="disable sha256 hashing of checkpoints to help loading performance", default=False) parser.add_argument("--no-download-sd-model", action='store_true', help="don't download SD1.5 model even if no model is found in --ckpt-dir", default=False) parser.add_argument("--profile", action='store_true', help="run profiler", default=False)