update import paths

cli/modules/bench.py

@@ -6,15 +6,10 @@ import asyncio
 import base64
 import io
 import json
-import os
-import sys
 import time
-
 from PIL import Image
-
-sys.path.append(os.path.join(os.path.dirname(__file__), 'modules'))
-import modules.sdapi as sdapi
-from modules.util import Map, log
+import sdapi as sdapi
+from util import Map, log
 
 
 options = Map({

cli/modules/embedding-preview.py

@@ -11,7 +11,6 @@ import argparse
 from pathlib import Path
 from PIL import Image
 from inspect import getsourcefile
-
 from util import Map, log
 from sdapi import getsync, postsync
 from grid import grid

cli/modules/extract-lora.py

@@ -0,0 +1,144 @@
+#!/bin/env python
+
+"""
+Extract approximating LoRA by SVD from two SD models
+Based on: <https://github.com/kohya-ss/sd-scripts/blob/main/networks/extract_lora_from_models.py>
+"""
+
+import os
+import sys
+import time
+import argparse
+import torch
+import transformers
+from tqdm import tqdm
+from util import log
+
+sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'modules', 'lora'))
+import library.model_util as model_util
+import networks.lora as lora
+
+
+def svd(args): # pylint: disable=redefined-outer-name
+    device = 'cuda' if torch.cuda.is_available() and args.device == 'cuda' else 'cpu'
+    transformers.logging.set_verbosity_error()
+    CLAMP_QUANTILE = 0.99
+    MIN_DIFF = 1e-6
+    if args.precision == 'fp32':
+        save_dtype = torch.float
+    elif args.precision == 'fp16':
+        save_dtype = torch.float16
+    elif args.precision == 'bf16':
+        save_dtype = torch.bfloat16
+    else:
+        save_dtype = None
+    t0 = time.time()
+    log.info({ 'loading model': args.original })
+    text_encoder_o, _, unet_o = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.original)
+    log.info({ 'loading model': args.tuned })
+    text_encoder_t, _, unet_t = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.tuned)
+    with torch.no_grad():
+        torch.cuda.empty_cache()
+    # create LoRA network to extract weights: Use dim (rank) as alpha
+    lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o)
+    lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t)
+    assert len(lora_network_o.text_encoder_loras) == len(lora_network_t.text_encoder_loras), 'model version is different'
+    # get diffs
+    diffs = {}
+    text_encoder_different = False
+    for i, (lora_o, lora_t) in enumerate(zip(lora_network_o.text_encoder_loras, lora_network_t.text_encoder_loras)):
+        lora_name = lora_o.lora_name
+        module_o = lora_o.org_module
+        module_t = lora_t.org_module
+        diff = module_t.weight - module_o.weight
+        # Text Encoder might be same
+        if torch.max(torch.abs(diff)) > MIN_DIFF:
+            text_encoder_different = True
+        diff = diff.float()
+        diffs[lora_name] = diff
+
+    if not text_encoder_different:
+        log.info({ 'lora': 'text encoder is same, extract U-Net only' })
+        lora_network_o.text_encoder_loras = []
+        diffs = {}
+
+    for i, (lora_o, lora_t) in enumerate(zip(lora_network_o.unet_loras, lora_network_t.unet_loras)):
+        lora_name = lora_o.lora_name
+        module_o = lora_o.org_module
+        module_t = lora_t.org_module
+        diff = module_t.weight - module_o.weight
+        diff = diff.float()
+        diff = diff.to(device)
+        diffs[lora_name] = diff
+    t1 = time.time()
+    log.info({ 'lora models': 'ready', 'time': round(t1 - t0, 2) })
+
+    # make LoRA with svd
+    log.info({ 'lora': 'calculating by svd' })
+    rank = args.dim
+    lora_weights = {}
+    with torch.no_grad():
+        for lora_name, mat in tqdm(list(diffs.items())):
+            conv2d = len(mat.size()) == 4
+            if conv2d:
+                mat = mat.squeeze()
+            U, S, Vh = torch.linalg.svd(mat)
+            U = U[:, :rank]
+            S = S[:rank]
+            U = U @ torch.diag(S)
+            Vh = Vh[:rank, :]
+            dist = torch.cat([U.flatten(), Vh.flatten()])
+            hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+            low_val = -hi_val
+            U = U.clamp(low_val, hi_val)
+            Vh = Vh.clamp(low_val, hi_val)
+            lora_weights[lora_name] = (U, Vh)
+    t2 = time.time()
+
+    # make state dict for LoRA
+    lora_network_o.apply_to(text_encoder_o, unet_o, text_encoder_different, True)     # to make state dict
+    lora_sd = lora_network_o.state_dict()
+    log.info({ 'lora extracted weights': len(lora_sd), 'time': round(t2 - t1, 2) })
+
+    for key in list(lora_sd.keys()):
+        if 'alpha' in key:
+            continue
+        lora_name = key.split('.')[0]
+        i = 0 if 'lora_up' in key else 1
+        weights = lora_weights[lora_name][i]
+        # print(key, i, weights.size(), lora_sd[key].size())
+        if len(lora_sd[key].size()) == 4: # pylint: disable=unsubscriptable-object
+            weights = weights.unsqueeze(2).unsqueeze(3)
+        assert weights.size() == lora_sd[key].size(), f'size unmatch: {key}' # pylint: disable=unsubscriptable-object
+        lora_sd[key] = weights  # pylint: disable=unsupported-assignment-operation
+
+    # load state dict to LoRA and save it
+    info = lora_network_o.load_state_dict(lora_sd)
+    log.info({ 'lora loading extracted weights': info })
+
+    dir_name = os.path.dirname(args.save)
+    if dir_name and not os.path.exists(dir_name):
+        os.makedirs(dir_name, exist_ok=True)
+
+    # minimum metadata
+    metadata = {'ss_network_dim': str(args.dim), 'ss_network_alpha': str(args.dim)}
+    lora_network_o.save_weights(args.save, save_dtype, metadata)
+    t3 = time.time()
+    log.info({ 'lora saved weights': args.save, 'time': round(t3 - t2, 2) })
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description = 'extract lora weights')
+    parser.add_argument('--v2', action='store_true', help='load Stable Diffusion v2.x model / Stable Diffusion')
+    parser.add_argument('--precision', type=str, default='fp16', choices=[None, 'fp32', 'fp16', 'bf16'], help='precision in saving, same to merging if omitted')
+    parser.add_argument('--device', type=str, default='cuda', choices=['cpu', 'cuda'], help='use cpu or cuda if available')
+    parser.add_argument('--original', type=str, default=None, required=True, help='Stable Diffusion original model: ckpt or safetensors file')
+    parser.add_argument('--tuned', type=str, default=None, required=True, help='Stable Diffusion tuned model, LoRA is difference of `original to tuned`: ckpt or safetensors file')
+    parser.add_argument('--save', type=str, default=None, required=True, help='destination file name: ckpt or safetensors file')
+    parser.add_argument('--dim', type=int, default=4, help='dimension (rank) of LoRA')
+    args = parser.parse_args()
+    log.info({ 'extract lora args': vars(args) })
+    if not os.path.exists(args.original) or not os.path.exists(args.tuned):
+        log.error({ 'models not found': [args.original, args.tuned] })
+    else:
+        svd(args)

cli/modules/grid.py

@@ -8,7 +8,6 @@ import argparse
 import math
 import logging
 from pathlib import Path
-
 import filetype
 from PIL import Image, ImageDraw, ImageFont
 from util import log

cli/modules/image-watermark.py

@@ -9,7 +9,7 @@ from imwatermark import WatermarkEncoder, WatermarkDecoder
 from PIL import Image
 from PIL.ExifTags import TAGS
 from PIL.TiffImagePlugin import ImageFileDirectory_v2
-from modules.util import log, Map
+from util import log, Map
 import piexif
 import piexif.helper
 

cli/modules/interrogate.py

@@ -8,10 +8,8 @@ import base64
 import sys
 import os
 import asyncio
-
 import filetype
 from PIL import Image
-
 from util import log, Map
 import sdapi as sdapi
 

cli/modules/models-preview.py

@@ -6,13 +6,12 @@ import time
 import asyncio
 import argparse
 from pathlib import Path
+from util import Map, log
+from sdapi import get, post, close
+from grid import grid
 
 sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
-sys.path.append(os.path.join(os.path.dirname(__file__), 'modules'))
 from generate import sd, generate
-from modules.util import Map, log
-from modules.sdapi import get, post, close
-from modules.grid import grid
 
 
 default = 'sd-v15-runwayml.ckpt [cc6cb27103]'

cli/modules/process.py

@@ -25,14 +25,12 @@ import base64
 import pathlib
 import argparse
 import logging
-
 import filetype
 import numpy as np
 import mediapipe as mp
 from PIL import Image, ImageOps
 from skimage.metrics import structural_similarity as ssim
 from scipy.stats import beta
-
 from util import log, Map
 from sdapi import postsync
 
@@ -40,6 +38,7 @@ from sdapi import postsync
 params = Map({
     'src': '', # source folder
     'dst': '', # destination folder
+    'format': '.png', # image format
     'extract_face': True, # extract face from image
     'extract_body': True, # extract face from image
     'clear_dst': True, # remove all files from destination at the start
@@ -50,15 +49,15 @@ params = Map({
     'face_pad': 0.07, # pad face image percentage
     'face_model': 1, # which face model to use 0/close-up 1/standard
     'face_blur_score': 1.5, # max score for face blur detection
-    'face_range_score': 0.3, # min score for face dynamic range detection
+    'face_range_score': 0.2, # min score for face dynamic range detection
     '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_score': 1.8, # max score for body blur detection
-    'body_range_score': 0.3, # min score for body dynamic range detection
-    'segmentation_face': True, # segmentation enabled
+    'body_range_score': 0.2, # min score for body dynamic range detection
+    'segmentation_face': False, # segmentation enabled
     'segmentation_body': False, # segmentation enabled
     'segmentation_model': 0, # segmentation model 0/general 1/landscape
     'segmentation_background': (192, 192, 192), # segmentation background color
@@ -263,7 +262,7 @@ def interrogate(img, fn):
     res = postsync('/sdapi/v1/interrogate', json)
     caption = res.caption if 'caption' in res else ''
     log.info({ 'interrogate': caption })
-    file = fn.replace('.jpg', '.txt')
+    file = fn.replace(params.format, '.txt')
     f = open(file, 'w')
     f.write(caption)
     f.close()
@@ -278,7 +277,7 @@ def process_file(f: str, dst: str = None, preview: bool = False, offline: bool =
         else:
             dir = dst
         base = os.path.basename(f).split('.')[0]
-        fn = os.path.join(dir, str(i[what]).rjust(3, '0') + '-' + what + '-' + base + '.jpg')
+        fn = os.path.join(dir, str(i[what]).rjust(3, '0') + '-' + what + '-' + base + params.format)
         # log.debug({ 'save': fn })
         if not preview:
             img.save(fn)
@@ -318,6 +317,7 @@ def process_file(f: str, dst: str = None, preview: bool = False, offline: bool =
         log.debug({ 'no body': f })
 
     image.close()
+    return i
 
 def process_images(src: str, dst: str, args = None):
     params.src = src
@@ -337,7 +337,7 @@ def process_images(src: str, dst: str, args = None):
         for root, _sub_dirs, files in os.walk(src):
             for f in files:
                 process_file(os.path.join(root, f), dst)
-
+    return i
 
 if __name__ == '__main__':
     # log.setLevel(logging.DEBUG)

cli/modules/prompt-ideas.py

@@ -6,9 +6,7 @@ model from: <https://huggingface.co/FredZhang7/distilgpt2-stable-diffusion-v2>
 
 import logging
 import argparse
-
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
-
 from util import log
 
 

cli/modules/prompt-promptist.py

@@ -5,9 +5,7 @@ use microsoft promptist to beautify prompt
 """
 
 import sys
-
 from transformers import AutoModelForCausalLM, AutoTokenizer
-
 from util import log
 
 

cli/modules/sdapi.py

@@ -9,7 +9,6 @@ import asyncio
 import logging
 import requests
 import sys
-
 from util import Map, log
 
 

cli/modules/train-lora.py

@@ -0,0 +1,186 @@
+#!/bin/env python
+
+"""
+Extract approximating LoRA by SVD from two SD models
+Based on: <https://github.com/kohya-ss/sd-scripts/blob/main/networks/train_network.py>
+"""
+
+import os
+import sys
+import argparse
+import tempfile
+import transformers
+from pathlib import Path
+from util import log, Map
+from process import process_file
+
+sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'modules', 'lora'))
+from train_network import train
+
+
+options = Map({
+    "v2": False,
+    "v_parameterization": False,
+    "pretrained_model_name_or_path": "/mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt",
+    "train_data_dir": "/tmp/rreid/img",
+    "shuffle_caption": False,
+    "caption_extension": ".txt",
+    "caption_extention": None,
+    "keep_tokens": None,
+    "color_aug": False,
+    "flip_aug": False,
+    "face_crop_aug_range": None,
+    "random_crop": False,
+    "debug_dataset": False,
+    "resolution": "512,512",
+    "cache_latents": True,
+    "enable_bucket": False,
+    "min_bucket_reso": 256,
+    "max_bucket_reso": 1024,
+    "bucket_reso_steps": 64,
+    "bucket_no_upscale": False,
+    "reg_data_dir": None,
+    "in_json": "/tmp/rreid/rreid.json",
+    "dataset_repeats": 1,
+    "output_dir": "/mnt/d/Models/lora/",
+    "output_name": "lora-rreid-random-v1",
+    "save_precision": "fp16",
+    "save_every_n_epochs": 1,
+    "save_n_epoch_ratio": None,
+    "save_last_n_epochs": None,
+    "save_last_n_epochs_state": None,
+    "save_state": False,
+    "resume": None,
+    "train_batch_size": 1,
+    "max_token_length": None,
+    "use_8bit_adam": False,
+    "mem_eff_attn": False,
+    "xformers": False,
+    "vae": None,
+    "learning_rate": 1e-05,
+    "max_train_steps": 5000,
+    "max_train_epochs": None,
+    "max_data_loader_n_workers": 8,
+    "persistent_data_loader_workers": False,
+    "seed": 42,
+    "gradient_checkpointing": False,
+    "gradient_accumulation_steps": 1,
+    "mixed_precision": "fp16",
+    "full_fp16": False,
+    "clip_skip": None,
+    "logging_dir": None,
+    "log_prefix": None,
+    "lr_scheduler": "cosine",
+    "lr_warmup_steps": 0,
+    "prior_loss_weight": 1.0,
+    "no_metadata": False,
+    "save_model_as": "ckpt",
+    "unet_lr": 0.001,
+    "text_encoder_lr": 5e-05,
+    "lr_scheduler_num_cycles": 1,
+    "lr_scheduler_power": 1,
+    "network_weights": None,
+    "network_module": "networks.lora",
+    "network_dim": 16,
+    "network_alpha": 1.0,
+    "network_args": None,
+    "network_train_unet_only": False,
+    "network_train_text_encoder_only": False,
+    "training_comment": "mood-magic"
+})
+
+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('--dir', type=str, default=None, required=True, help='folder containing lora checkpoints')
+    parser.add_argument('--name', type=str, default=None, required=True, help='lora name')
+    parser.add_argument('--steps', type=int, default=5000, required=False, help='training steps')
+    parser.add_argument('--dim', type=int, default=16, required=False, help='network dimension')
+    parser.add_argument("--noprocess", default = False, action='store_true', help = "skip processing and use existing input data")
+    args = parser.parse_args()
+    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.output_dir = args.dir
+    options.output_name = args.name
+    options.max_train_steps = args.steps
+    options.network_dim = args.dim
+    log.info({ 'train lora args': vars(options) })
+    transformers.logging.set_verbosity_error()
+
+    if args.noprocess:
+        options.train_data_dir = args.input
+    else:
+        dir = os.path.join(tempfile.gettempdir(), args.name, '10_processed')
+        Path(dir).mkdir(parents=True, exist_ok=True)
+        files = []
+        json_data = {}
+        for root, _sub_dirs, folder in os.walk(args.input):
+            for f in folder:
+                files.append(os.path.join(root, f))
+        for f in files:
+            res = process_file(f = f, dst = dir, preview = False, offline = True)
+            
+        log.info({ 'processed': res, 'inputs': len(files) })
+        options.train_data_dir = args.input
+        dir = os.path.join(tempfile.gettempdir(), args.name)
+
+    train(options)
+
+
+"""
+- cannot use `accelerate` with *dynamo* enabled
+- cannot use `xformers` due to *faketensors* requirement
+- cannot use `mem_eff_attn` due to *forwardfunc* mismatch
+
+TODO
+
+--gradient_checkpointing
+--gradient_accumulation_steps=10
+--caption_extension=txt
+--in_json
+
+WORKING
+
+process.py --output "/tmp/rreid/img/10_processed" /home/vlado/generative/Input/ryanreid/random --offline
+
+accelerate launch --no_python --quiet --num_cpu_threads_per_process=16 python /home/vlado/dev/automatic/modules/lora/train_network.py \
+--pretrained_model_name_or_path="/mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt" \
+--train_data_dir="/tmp/rreid/img" \
+--logging_dir="/tmp/rreid/logging" \
+--output_dir="/mnt/d/Models/lora/" \
+--output_name="lora-rreid-random-v1" \
+--resolution=512,512 \
+--learning_rate=1e-5 \
+--unet_lr=1e-3 \
+--text_encoder_lr=5e-5 \
+--lr_scheduler_num_cycles=1 \
+--lr_scheduler=cosine \
+--max_train_steps=5000 \
+--network_alpha=1 \
+--network_dim=16 \
+--network_module=networks.lora \
+--save_every_n_epochs=1 \
+--save_model_as=ckpt \
+--save_precision=fp16 \
+--mixed_precision=fp16 \
+--seed=42 \
+--train_batch_size=1 \
+--cache_latents \
+
+metadata { image_key: img_md: { caption: str, tags: [] } }
+
+abs_path = glob_images(train_data_dir, image_key)
+
+}}
+
+./train-lora.py --model /mnt/d/Models/stable-diffusion/sd-v15-runwayml.ckpt --name rreid --dir /mnt/d/Models/lora --input ~/generative/Input/ryanreid/random/
+"""

cli/modules/train-losschart.py

@@ -6,12 +6,10 @@ import sys
 import json
 import pathlib
 import logging
-
 import torch
 import numpy as np
 from PIL import Image, ImageFont, ImageDraw
 from matplotlib import pyplot as plt
-
 from util import log, Map
 
 

cli/modules/train-lossrate.py

@@ -5,11 +5,9 @@ auto-generate learn-rate
 import io
 import math
 import logging
-
 import numpy as np
 from PIL import Image, ImageFont, ImageDraw
 from matplotlib import pyplot as plt
-
 from util import log, Map
 
 

cli/modules/video-extract.py

@@ -8,7 +8,6 @@ import subprocess
 import pathlib
 import argparse
 import filetype
-
 from util import log, Map