add sdp cross-optimization

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
 

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)

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())

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

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

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_)

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)