lora with multi-concept training

2026-09-19 01:04:32 +02:00 · 2023-03-07 13:29:46 -05:00
parent 809b793542
commit ce2eb4e272
8 changed files with 126 additions and 106 deletions
@@ -71,7 +71,7 @@ async def interrogate(f):
    keywords = {}
    if 'caption' in res:
        for term in res.caption.split(', '):
-            term = term.replace('(', '').replace(')', '').split(':')
+            term = term.replace('(', '').replace(')', '').replace('\\', '').split(':')
            if len(term) < 2:
                continue
            keywords[term[0]] = term[1]
@@ -342,7 +342,7 @@ def interrogate(img, fn, intag = None):
        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]
+            tags = [t.replace('(', '').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())
@@ -385,22 +385,16 @@ def process_file(f: str, dst: str = None, preview: bool = False, offline: bool =
        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
+    if len(opts) > 0:
+        params.keep_original = True if 'original' in opts else False
+        params.extract_face = True if 'face' in opts else False
+        params.extract_body = True if 'body' in opts else False
+        params.face_blur = True if 'blur' in opts else False
+        params.body_blur = True if 'blur' in opts else False
+        params.face_range = True if 'range' in opts else False
+        params.body_range = True if 'range' in opts else False
+        params.face_upscale = True if 'upscale' in opts else False
+        params.face_restore = True if 'upscale' in opts else False

    log.info({ 'processing': f })
    try:
@@ -19,6 +19,7 @@ import gc
 import sys
 import json
 import time
+import shutil
 import argparse
 import tempfile
 import torch
@@ -39,78 +40,78 @@ sys.path.append(locon_path)
 from train_network import train

 options = Map({
-    "v2": False,
-    "v_parameterization": False,
-    "pretrained_model_name_or_path": "",
-    "train_data_dir": "",
-    "shuffle_caption": False,
+    "bucket_no_upscale": False,
+    "bucket_reso_steps": 64,
+    "cache_latents": True,
+    "caption_dropout_every_n_epochs": None,
+    "caption_dropout_rate": 0.0,
    "caption_extension": ".txt",
    "caption_extention": ".txt",
-    "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": "",
-    "dataset_repeats": 1,
-    "output_dir": "",
-    "output_name": "",
-    "save_precision": "fp16",
-    "save_every_n_epochs": None,
-    "save_n_epoch_ratio": None,
-    "save_last_n_epochs": None,
-    "save_last_n_epochs_state": None,
-    "save_state": False,
-    "resume": None,
-    "max_grad_norm": 0.0,
-    "train_batch_size": 1,
-    "max_token_length": None,
-    "use_8bit_adam": False,
-    "mem_eff_attn": False,
-    "xformers": False,
-    "vae": None,
-    "learning_rate": 1e-04,
-    "max_train_steps": 8000,
-    "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,
+    "caption_tag_dropout_rate": 0.0,
    "clip_skip": None,
-    "logging_dir": None,
+    "color_aug": False,
+    "dataset_repeats": 1,
+    "debug_dataset": False,
+    "enable_bucket": False,
+    "face_crop_aug_range": None,
+    "flip_aug": False,
+    "full_fp16": False,
+    "gradient_accumulation_steps": 1,
+    "gradient_checkpointing": False,
+    "in_json": "",
+    "keep_tokens": None,
+    "learning_rate": 5e-05,
    "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,
+    "logging_dir": None,
    "lr_scheduler_num_cycles": 1,
    "lr_scheduler_power": 1,
-    "network_weights": None,
-    "network_module": "networks.lora",
-    "network_dim": 16,
+    "lr_scheduler": "cosine",
+    "lr_warmup_steps": 0,
+    "max_bucket_reso": 1024,
+    "max_data_loader_n_workers": 8,
+    "max_grad_norm": 0.0,
+    "max_token_length": None,
+    "max_train_epochs": None,
+    "max_train_steps": 5000,
+    "mem_eff_attn": False,
+    "min_bucket_reso": 256,
+    "mixed_precision": "fp16",
    "network_alpha": 1.0,
    "network_args": None,
-    "network_train_unet_only": False,
+    "network_dim": 16,
+    "network_module": "networks.lora",
    "network_train_text_encoder_only": False,
+    "network_train_unet_only": False,
+    "network_weights": None,
+    "no_metadata": False,
+    "output_dir": "",
+    "output_name": "",
+    "persistent_data_loader_workers": False,
+    "pretrained_model_name_or_path": "",
+    "prior_loss_weight": 1.0,
+    "random_crop": False,
+    "reg_data_dir": None,
+    "resolution": "512,512",
+    "resume": None,
+    "save_every_n_epochs": None,
+    "save_last_n_epochs_state": None,
+    "save_last_n_epochs": None,
+    "save_model_as": "ckpt",
+    "save_n_epoch_ratio": None,
+    "save_precision": "fp16",
+    "save_state": False,
+    "seed": 42,
+    "shuffle_caption": False,
+    "text_encoder_lr": 5e-05,
+    "train_batch_size": 1,
+    "train_data_dir": "",
    "training_comment": "mood-magic",
-    "caption_dropout_rate": 0.0,
-    "caption_dropout_every_n_epochs": None,
-    "caption_tag_dropout_rate": 0.0,
+    "unet_lr": 0.001,
+    "use_8bit_adam": False,
+    "v_parameterization": False,
+    "v2": False,
+    "vae": None,
+    "xformers": False,
 })


@@ -134,7 +135,7 @@ if __name__ == '__main__':
    parser.add_argument('--tag', type=str, default=None, required=False, help='primary tag')
    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('--process', type=str, default='original', required=False, 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')
@@ -145,7 +146,7 @@ if __name__ == '__main__':
    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('--alpha', type=float, default=0, required=False, help='alpha for weights scaling, default: half of dim')
    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')
@@ -182,6 +183,7 @@ if __name__ == '__main__':
    options.output_name = args.output
    options.max_train_steps = args.steps
    options.network_dim = args.dim
+    options.network_alpha = args.dim // 2 if args.alpha == 0 else args.alpha
    options.gradient_accumulation_steps = args.gradient
    options.save_every_n_epochs = args.interim if args.interim > 0 else None
    options.learning_rate = args.lr
@@ -193,13 +195,9 @@ if __name__ == '__main__':
    transformers.logging.set_verbosity_error()
    mem_stats()

-    concept = 'lora'
-    if args.tag is not None:
-        concept = args.tag.split(',')[0].strip()
-    dir = os.path.join(tempfile.gettempdir(), args.output, str(args.repeats) + '_' + concept)
-    Path(dir).mkdir(parents=True, exist_ok=True)
    json_file = os.path.join(tempfile.gettempdir(), args.output, args.output + '.json')
-    options.train_data_dir = os.path.join(tempfile.gettempdir(), args.output)
+    base = os.path.join(tempfile.gettempdir(), args.output)
+    options.train_data_dir = base
    res = None

    if args.dreambooth:
@@ -215,28 +213,52 @@ if __name__ == '__main__':
        # preprocess
        processing_options = args.process.split(',')
        processing_options = [opt.strip() for opt in re.split(',| ', args.process)]
-        log.info({ 'processing options': processing_options })
+        log.info({ 'processing steps': 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, tag = args.tag, opts = processing_options)
-                if not args.dreambooth:
-                    with open(json_file, "w") as outfile:
-                        outfile.write(json.dumps(metadata, indent=2))
-            except ValueError as e:
-                exit(1)
+
+        steps = [step for step in processing_options if step in ['face', 'body', 'original']]
+        for step in steps:
+            # processing_options = [step for step in processing_options if step not in ['face', 'body', 'original']].append(step)
+            if step == 'face':
+                opts = [step for step in processing_options if step not in ['body', 'original']]
+            if step == 'body':
+                opts = [step for step in processing_options if step not in ['face', 'original', 'upscale', 'restore']]
+            if step == 'original':
+                opts = [step for step in processing_options if step not in ['face', 'body', 'upscale', 'restore', 'blur', 'range']]
+            log.info({ 'processing step': opts })
+            concept = step
+            if concept == 'original' and args.tag is not None:
+                concept = args.tag.split(',')[0].strip()
+            dir = os.path.join(base, str(args.repeats) + '_' + concept)
+            if os.path.exists(dir):
+                shutil.rmtree(dir, ignore_errors=True)
+            Path(dir).mkdir(parents=True, exist_ok=True)
+
+            for f in files:
+                try:
+                    res, metadata = modules.process.process_file(f = f, dst = dir, preview = False, offline = args.offline, txt = args.dreambooth, tag = args.tag, opts = opts)
+                    if not args.dreambooth:
+                        with open(json_file, "w") as outfile:
+                            outfile.write(json.dumps(metadata, indent=2))
+                except ValueError as e:
+                    exit(1)
+            log.info({ 'processed step': step, 'outputs': res, 'inputs': len(files), 'metadata': json_file, 'path': dir })
+
        modules.process.unload_models()
        mem_stats()

+
+    dirs = [os.path.join(base, dir) for dir in os.listdir(base) if os.path.isdir(os.path.join(base, dir))]
+    log.info({ 'input datasets': dirs, 'metadata': json_file })
+
    if not args.nolatents and not args.dreambooth:
        # create latents
-        latents.create_vae_latents(Map({ 'input': dir, 'json': json_file }))
-        latents.unload_vae()
+        for dir in dirs:
+            latents.create_vae_latents(Map({ 'input': dir, 'json': json_file }))
+            latents.unload_vae()
        mem_stats()
-
-        log.info({ 'processed': res, 'inputs': len(files), 'metadata': json_file, 'path': dir })
    else:
        log.info({ 'skip processing': len(files), 'metadata': json_file, 'path': dir })