3 files changed, 73 insertions, 57 deletions

diff --git a/infer.py b/infer.py
index ed86ab1..93848d7 100644
--- a/infer.py
+++ b/infer.py
@@ -26,6 +26,8 @@ from diffusers import (
     DEISMultistepScheduler,
     UniPCMultistepScheduler
 )
+from peft import LoraConfig, LoraModel, set_peft_model_state_dict
+from safetensors.torch import load_file
 from transformers import CLIPTextModel
 
 from data.keywords import str_to_keywords, keywords_to_str
@@ -43,7 +45,7 @@ default_args = {
     "model": "stabilityai/stable-diffusion-2-1",
     "precision": "fp32",
     "ti_embeddings_dir": "embeddings_ti",
-    "lora_embeddings_dir": "embeddings_lora",
+    "lora_embedding": None,
     "output_dir": "output/inference",
     "config": None,
 }
@@ -99,7 +101,7 @@ def create_args_parser():
         type=str,
     )
     parser.add_argument(
-        "--lora_embeddings_dir",
+        "--lora_embedding",
         type=str,
     )
     parser.add_argument(
@@ -236,6 +238,38 @@ def load_embeddings(pipeline, embeddings_dir):
     print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
 
 
+def load_lora(pipeline, path):
+    if path is None:
+        return
+
+    path = Path(path)
+
+    with open(path / "lora_config.json", "r") as f:
+        lora_config = json.load(f)
+
+    tensor_files = list(path.glob("*_end.safetensors"))
+
+    if len(tensor_files) == 0:
+        return
+
+    lora_checkpoint_sd = load_file(path / tensor_files[0])
+    unet_lora_ds = {k: v for k, v in lora_checkpoint_sd.items() if "text_encoder_" not in k}
+    text_encoder_lora_ds = {
+        k.replace("text_encoder_", ""): v for k, v in lora_checkpoint_sd.items() if "text_encoder_" in k
+    }
+
+    unet_config = LoraConfig(**lora_config["peft_config"])
+    pipeline.unet = LoraModel(unet_config, pipeline.unet)
+    set_peft_model_state_dict(pipeline.unet, unet_lora_ds)
+
+    if "text_encoder_peft_config" in lora_config:
+        text_encoder_config = LoraConfig(**lora_config["text_encoder_peft_config"])
+        pipeline.text_encoder = LoraModel(text_encoder_config, pipeline.text_encoder)
+        set_peft_model_state_dict(pipeline.text_encoder, text_encoder_lora_ds)
+
+    return
+
+
 def create_scheduler(config, scheduler: str, subscheduler: Optional[str] = None):
     if scheduler == "plms":
         return PNDMScheduler.from_config(config)
@@ -441,6 +475,7 @@ def main():
     pipeline = create_pipeline(args.model, dtype)
 
     load_embeddings(pipeline, args.ti_embeddings_dir)
+    load_lora(pipeline, args.lora_embedding)
     # pipeline.unet.load_attn_procs(args.lora_embeddings_dir)
 
     cmd_parser = create_cmd_parser()
diff --git a/train_lora.py b/train_lora.py
index 73b3e19..1ca56d9 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -1,7 +1,6 @@
 import argparse
 import datetime
 import logging
-import itertools
 from pathlib import Path
 from functools import partial
 import math
@@ -247,9 +246,15 @@ def parse_args():
         help="Automatically find a learning rate (no training).",
     )
     parser.add_argument(
-        "--learning_rate",
+        "--learning_rate_unet",
         type=float,
-        default=2e-6,
+        default=1e-4,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--learning_rate_text",
+        type=float,
+        default=5e-5,
         help="Initial learning rate (after the potential warmup period) to use.",
     )
     parser.add_argument(
@@ -548,13 +553,18 @@ def main():
         print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
 
     if args.scale_lr:
-        args.learning_rate = (
-            args.learning_rate * args.gradient_accumulation_steps *
+        args.learning_rate_unet = (
+            args.learning_rate_unet * args.gradient_accumulation_steps *
+            args.train_batch_size * accelerator.num_processes
+        )
+        args.learning_rate_text = (
+            args.learning_rate_text * args.gradient_accumulation_steps *
             args.train_batch_size * accelerator.num_processes
         )
 
     if args.find_lr:
-        args.learning_rate = 1e-6
+        args.learning_rate_unet = 1e-6
+        args.learning_rate_text = 1e-6
         args.lr_scheduler = "exponential_growth"
 
     if args.optimizer == 'adam8bit':
@@ -611,8 +621,8 @@ def main():
         )
 
         args.lr_scheduler = "adafactor"
-        args.lr_min_lr = args.learning_rate
-        args.learning_rate = None
+        args.lr_min_lr = args.learning_rate_unet
+        args.learning_rate_unet = None
     elif args.optimizer == 'dadam':
         try:
             import dadaptation
@@ -628,7 +638,8 @@ def main():
             d0=args.dadaptation_d0,
         )
 
-        args.learning_rate = 1.0
+        args.learning_rate_unet = 1.0
+        args.learning_rate_text = 1.0
     elif args.optimizer == 'dadan':
         try:
             import dadaptation
@@ -642,7 +653,8 @@ def main():
             d0=args.dadaptation_d0,
         )
 
-        args.learning_rate = 1.0
+        args.learning_rate_unet = 1.0
+        args.learning_rate_text = 1.0
     else:
         raise ValueError(f"Unknown --optimizer \"{args.optimizer}\"")
 
@@ -695,15 +707,16 @@ def main():
         sample_frequency = math.ceil(num_train_epochs * (sample_frequency / args.num_train_steps))
 
     optimizer = create_optimizer(
-        (
-            param
-            for param in itertools.chain(
-                unet.parameters(),
-                text_encoder.parameters(),
-            )
-            if param.requires_grad
-        ),
-        lr=args.learning_rate,
+        [
+            {
+                "params": unet.parameters(),
+                "lr": args.learning_rate_unet,
+            },
+            {
+                "params": text_encoder.parameters(),
+                "lr": args.learning_rate_text,
+            },
+        ]
     )
 
     lr_scheduler = get_scheduler(
diff --git a/training/functional.py b/training/functional.py
index 06848cb..c30d1c0 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -321,45 +321,13 @@ def loss_step(
     )
 
     if offset_noise_strength != 0:
-        solid_image = partial(
-            make_solid_image,
-            shape=images.shape[1:],
-            vae=vae,
+        offset_noise = torch.randn(
+            (latents.shape[0], latents.shape[1], 1, 1),
             dtype=latents.dtype,
             device=latents.device,
             generator=generator
-        )
-
-        white_cache_key = f"img_white_{images.shape[2]}_{images.shape[3]}"
-        black_cache_key = f"img_black_{images.shape[2]}_{images.shape[3]}"
-
-        if white_cache_key not in cache:
-            img_white = solid_image(1)
-            cache[white_cache_key] = img_white
-        else:
-            img_white = cache[white_cache_key]
-
-        if black_cache_key not in cache:
-            img_black = solid_image(0)
-            cache[black_cache_key] = img_black
-        else:
-            img_black = cache[black_cache_key]
-
-        offset_strength = torch.rand(
-            (bsz, 1, 1, 1),
-            dtype=latents.dtype,
-            layout=latents.layout,
-            device=latents.device,
-            generator=generator
-        )
-        offset_strength = offset_noise_strength * (offset_strength * 2 - 1)
-        offset_images = torch.where(
-            offset_strength >= 0,
-            img_white.expand(noise.shape),
-            img_black.expand(noise.shape)
-        )
-        offset_strength = offset_strength.abs().expand(noise.shape)
-        noise = slerp(noise, offset_images, offset_strength, zdim=(-1, -2))
+        ).expand(noise.shape)
+        noise += offset_noise_strength * offset_noise
 
     # Sample a random timestep for each image
     timesteps = torch.randint(