integrate locon

README.md

@@ -6,19 +6,34 @@
 
 <br>
 
-## Notes
+### Notes
 
 Fork is as close as up-to-date with origin as time allows  
 All code changes are merged upstream whenever possible  
 
 Fork adds extra functionality:
 - New skin and UI layout  
-- Ships with additional **extensions**  
-  e.g. `System Info`, `Steps Animation`, etc.  
 - Ships with set of **CLI** tools that rely on *SD API* for execution:  
   e.g. `generate`, `train`, `bench`, etc.  
   [Full list](<cli/>)
 
+### Integrated Extensions:
+
+- [System Info](https://github.com/vladmandic/sd-extension-system-info)
+- [ControlNet](https://github.com/Mikubill/sd-webui-controlnet)
+- [Image Browser](https://github.com/AlUlkesh/stable-diffusion-webui-images-browser)
+- [LORA](https://github.com/kohya-ss/sd-scripts) (both training and inference)
+- [LoCon](https://github.com/KohakuBlueleaf/LoCon) (both training and inference)
+- [Model Converter](https://github.com/Akegarasu/sd-webui-model-converter)
+- [CLiP Interrogator](https://github.com/pharmapsychotic/clip-interrogator-ext)
+- [Dynamic Thresholding](https://github.com/mcmonkeyprojects/sd-dynamic-thresholding)
+- [Steps Animation](https://github.com/vladmandic/sd-extension-steps-animation)
+- [Seed Travel](https://github.com/yownas/seed_travel)
+
+*Note*: Extensions are automatically updated to latest version on `install`
+
+### Start Script
+
 Simplified start script: `automatic.sh`  
 *Existing `webui.sh`/`webui.bat` scripts still exist for backward compatibility*  
 
@@ -62,7 +77,7 @@ Recommended to run `install` after `update`
 
 ## Install
 
-1. Install `PyTorch` first
+1. Install `Python`, `Git` and `PyTorch` first
 2. Clone and initialize repository
 
 > git clone https://github.com/vladmandic/automatic  
@@ -82,13 +97,14 @@ Recommended to run `install` after `update`
       Detached repos
       Local changes
  
-
 <br>
 
 ## Differences
 
 Fork does differ in few things:
 - Drops compatibility with `python` **3.7** and requires **3.9**  
+  Recommended is **Python 3.10**  
+  Note that **Python 3.11** or **3.12** are NOT supported  
 - Updated **Python** libraries to latest known compatible versions  
   e.g. `accelerate`, `transformers`, `numpy`, etc.  
 - Includes opinionated **System** and **Options** configuration  

cli/modules/models-diff.py

@@ -0,0 +1,74 @@
+#!/bin/env python
+# based on <https://huggingface.co/JosephusCheung/ASimilarityCalculatior>
+
+import safetensors
+import sys
+import torch
+from pathlib import Path
+import torch.nn as nn
+import torch.nn.functional as F
+import warnings
+from util import log
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+def cal_cross_attn(to_q, to_k, to_v, rand_input):
+    hidden_dim, embed_dim = to_q.shape
+    attn_to_q = nn.Linear(hidden_dim, embed_dim, bias=False)
+    attn_to_k = nn.Linear(hidden_dim, embed_dim, bias=False)
+    attn_to_v = nn.Linear(hidden_dim, embed_dim, bias=False)
+    attn_to_q.load_state_dict({"weight": to_q})
+    attn_to_k.load_state_dict({"weight": to_k})
+    attn_to_v.load_state_dict({"weight": to_v})
+    
+    return torch.einsum(
+        "ik, jk -> ik", 
+        F.softmax(torch.einsum("ij, kj -> ik", attn_to_q(rand_input), attn_to_k(rand_input)), dim=-1),
+        attn_to_v(rand_input)
+    )
+       
+def load_model(path):
+    if path.suffix == ".safetensors":
+        return safetensors.torch.load_file(path, device="cpu")
+    else:
+        ckpt = torch.load(path, map_location="cpu")
+        return ckpt["state_dict"] if "state_dict" in ckpt else ckpt
+        
+def eval(model, n, input):
+    qk = f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_q.weight"
+    uk = f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_k.weight"
+    vk = f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_v.weight"
+    atoq, atok, atov = model[qk], model[uk], model[vk]
+    attn = cal_cross_attn(atoq, atok, atov, input)
+    return attn
+
+def main():
+    file1 = Path(sys.argv[1])
+    files = sys.argv[2:]
+    seed = 114514
+    torch.manual_seed(seed)
+    model_a = load_model(file1)   
+    log.info(f"base: {file1.name}")
+
+    map_attn_a = {}
+    map_rand_input = {}
+    for n in range(3, 11):
+        hidden_dim, embed_dim = model_a[f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_q.weight"].shape
+        rand_input = torch.randn([embed_dim, hidden_dim])
+        map_attn_a[n] = eval(model_a, n, rand_input)
+        map_rand_input[n] = rand_input       
+    del model_a
+     
+    for file2 in files:
+        file2 = Path(file2)
+        model_b = load_model(file2)
+        sims = []
+        for n in range(3, 11):
+            attn_a = map_attn_a[n]
+            attn_b = eval(model_b, n, map_rand_input[n])
+            sim = torch.mean(torch.cosine_similarity(attn_a, attn_b))
+            sims.append(sim)
+        log.info(f"{file2}: {torch.mean(torch.stack(sims)) * 1e2:.2f}%")
+        
+if __name__ == "__main__":
+    main()

cli/train-lora.py

@@ -31,10 +31,12 @@ import modules.sdapi
 
 latents = importlib.import_module('modules.lora-latents')
 
-sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'modules', 'lora'))
+lora_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'lora'))
+sys.path.append(lora_path)
+locon_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'locon'))
+sys.path.append(locon_path)
 from train_network import train
 
-
 options = Map({
     "v2": False,
     "v_parameterization": False,
@@ -138,14 +140,16 @@ if __name__ == '__main__':
     parser.add_argument('--offline', default = False, action='store_true', help = 'do not use webui server for processing')
     parser.add_argument('--shutdown', default = False, action='store_true', help = 'shutdown webui server')
     parser.add_argument('--gradient', type=int, default=1, required=False, help='gradient accumulation steps, default: %(default)s')
-    parser.add_argument('--steps', type=int, default=5000, required=False, help='training steps, default: %(default)s')
-    parser.add_argument('--dim', type=int, default=128, required=False, help='network dimension, default: %(default)s')
+    parser.add_argument('--steps', type=int, default=4000, required=False, help='training steps, default: %(default)s')
+    parser.add_argument('--dim', type=int, default=40, required=False, help='network dimension, default: %(default)s')
     parser.add_argument('--repeats', type=int, default=10, required=False, help='number of repeats per image, default: %(default)s')
+    parser.add_argument('--alpha', type=float, default=1, required=False, help='alpha for weights scaling, default: %(default)s')
     parser.add_argument('--batch', type=int, default=1, required=False, help='batch size, default: %(default)s')
     parser.add_argument('--lr', type=float, default=1e-04, required=False, help='model learning rate, default: %(default)s')
     parser.add_argument('--unetlr', type=float, default=1e-04, required=False, help='unet learning rate, default: %(default)s')
     parser.add_argument('--textlr', type=float, default=5e-05, required=False, help='text encoder learning rate, default: %(default)s')
     parser.add_argument('--dreambooth', default=False, action='store_true', help = "use dreambooth style training")
+    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()
     if args.debug:
@@ -171,6 +175,7 @@ if __name__ == '__main__':
     options.unet_lr = args.unetlr
     options.text_encoder_lr = args.textlr
     options.train_batch_size = args.batch
+    options.network_alpha = args.alpha
     log.info({ 'train lora args': vars(options) })
     transformers.logging.set_verbosity_error()
     mem_stats()
@@ -185,6 +190,7 @@ if __name__ == '__main__':
     res = None
 
     if args.dreambooth:
+        log.info({ 'using dreambooth style training': True })
         options.in_json = None
     else:
         options.in_json = json_file
@@ -226,6 +232,13 @@ if __name__ == '__main__':
         modules.sdapi.shutdown()
         time.sleep(1)
 
+    if args.locon:
+        # python3 sd-scripts/train_network.py --network_module locon.locon_kohya --network_dim "RANK_FOR_TRANSFORMER" --network_alpha "ALPHA_FOR_TRANSFORMER" --network_args "conv_dim=RANK_FOR_CONV" "conv_alpha=ALPHA_FOR_CONV" "dropout=DROPOUT_RATE"
+        # options.network_dim = 'RANK_FOR_TRANSFORMER'
+        # options.network_alpha = 'ALPHA_FOR_TRANSFORMER'
+        # options.network_args = ['conv_dim=RANK_FOR_CONV', 'conv_alpha=ALPHA_FOR_CONV', 'dropout=DROPOUT_RATE']
+        log.info({ 'using locon network': True })
+        options.network_module = 'locon.locon_kohya'
     if not args.notrain:
         train(options)
         mem_stats()

ui-config.json

@@ -1316,8 +1316,24 @@
     "customscript/seed_travel.py/img2img/SSIM threshold (1.0=exact copy)/minimum": 0.0,
     "customscript/seed_travel.py/img2img/SSIM threshold (1.0=exact copy)/maximum": 1.0,
     "customscript/seed_travel.py/img2img/SSIM threshold (1.0=exact copy)/step": 0.01,
-    "customscript/posex.py/txt2img/Send this image to ControlNet./visible": true,
-    "customscript/posex.py/txt2img/Send this image to ControlNet./value": false,
-    "customscript/posex.py/img2img/Send this image to ControlNet./visible": true,
-    "customscript/posex.py/img2img/Send this image to ControlNet./value": false
+    "customscript/seed_travel.py/txt2img/SSIM threshold (0 to disable)/visible": true,
+    "customscript/seed_travel.py/txt2img/SSIM threshold (0 to disable)/value": 0.0,
+    "customscript/seed_travel.py/txt2img/SSIM threshold (0 to disable)/minimum": 0.0,
+    "customscript/seed_travel.py/txt2img/SSIM threshold (0 to disable)/maximum": 1.0,
+    "customscript/seed_travel.py/txt2img/SSIM threshold (0 to disable)/step": 0.01,
+    "customscript/seed_travel.py/txt2img/SSIM CenterCrop% (0 to disable)/visible": true,
+    "customscript/seed_travel.py/txt2img/SSIM CenterCrop% (0 to disable)/value": 0,
+    "customscript/seed_travel.py/txt2img/SSIM CenterCrop% (0 to disable)/minimum": 0,
+    "customscript/seed_travel.py/txt2img/SSIM CenterCrop% (0 to disable)/maximum": 100,
+    "customscript/seed_travel.py/txt2img/SSIM CenterCrop% (0 to disable)/step": 1,
+    "customscript/seed_travel.py/img2img/SSIM threshold (0 to disable)/visible": true,
+    "customscript/seed_travel.py/img2img/SSIM threshold (0 to disable)/value": 0.0,
+    "customscript/seed_travel.py/img2img/SSIM threshold (0 to disable)/minimum": 0.0,
+    "customscript/seed_travel.py/img2img/SSIM threshold (0 to disable)/maximum": 1.0,
+    "customscript/seed_travel.py/img2img/SSIM threshold (0 to disable)/step": 0.01,
+    "customscript/seed_travel.py/img2img/SSIM CenterCrop% (0 to disable)/visible": true,
+    "customscript/seed_travel.py/img2img/SSIM CenterCrop% (0 to disable)/value": 0,
+    "customscript/seed_travel.py/img2img/SSIM CenterCrop% (0 to disable)/minimum": 0,
+    "customscript/seed_travel.py/img2img/SSIM CenterCrop% (0 to disable)/maximum": 100,
+    "customscript/seed_travel.py/img2img/SSIM CenterCrop% (0 to disable)/step": 1
 }