2 files changed, 84 insertions, 21 deletions

diff --git a/training/functional.py b/training/functional.py
index 8dc2b9f..43ee356 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -251,6 +251,25 @@ def add_placeholder_tokens(
     return placeholder_token_ids, initializer_token_ids
 
 
+def snr_weight(noisy_latents, latents, gamma):
+    if gamma:
+        sigma = torch.sub(noisy_latents, latents)
+        zeros = torch.zeros_like(sigma)
+        alpha_mean_sq = F.mse_loss(latents.float(), zeros.float(), reduction="none").mean([1, 2, 3])
+        sigma_mean_sq = F.mse_loss(sigma.float(), zeros.float(), reduction="none").mean([1, 2, 3])
+        snr = torch.div(alpha_mean_sq, sigma_mean_sq)
+        gamma_over_snr = torch.div(torch.ones_like(snr) * gamma, snr)
+        snr_weight = torch.minimum(gamma_over_snr, torch.ones_like(gamma_over_snr)).float()
+        return snr_weight
+
+    return torch.tensor(
+        [1],
+        dtype=latents.dtype,
+        layout=latents.layout,
+        device=latents.device,
+    )
+
+
 def loss_step(
     vae: AutoencoderKL,
     noise_scheduler: SchedulerMixin,
@@ -308,21 +327,13 @@ def loss_step(
     model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
 
     # Get the target for loss depending on the prediction type
-    alpha_t = noise_scheduler.alphas_cumprod[timesteps].float()
-    snr = alpha_t / (1 - alpha_t)
-    min_snr = snr.clamp(max=min_snr_gamma)
-
     if noise_scheduler.config.prediction_type == "epsilon":
         target = noise
-        loss_weight = min_snr / snr
     elif noise_scheduler.config.prediction_type == "v_prediction":
         target = noise_scheduler.get_velocity(latents, noise, timesteps)
-        loss_weight = min_snr / (snr + 1)
     else:
         raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
 
-    loss_weight = loss_weight[..., None, None, None]
-
     if with_prior_preservation:
         # Chunk the noise and model_pred into two parts and compute the loss on each part separately.
         model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0)
@@ -339,7 +350,11 @@ def loss_step(
     else:
         loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
 
-    loss = (loss_weight * loss).mean([1, 2, 3]).mean()
+    loss = loss.mean([1, 2, 3])
+
+    loss_weight = snr_weight(noisy_latents, latents, min_snr_gamma)
+    loss = (loss_weight * loss).mean()
+
     acc = (model_pred == target).float().mean()
 
     return loss, acc, bsz
@@ -412,7 +427,7 @@ def train_loop(
     try:
         for epoch in range(num_epochs):
             if accelerator.is_main_process:
-                if epoch % sample_frequency == 0 and epoch != 0:
+                if epoch % sample_frequency == 0:
                     local_progress_bar.clear()
                     global_progress_bar.clear()
 
diff --git a/training/strategy/lora.py b/training/strategy/lora.py
index cab5e4c..aa75bec 100644
--- a/training/strategy/lora.py
+++ b/training/strategy/lora.py
@@ -2,6 +2,7 @@ from typing import Optional
 from functools import partial
 from contextlib import contextmanager
 from pathlib import Path
+import itertools
 
 import torch
 from torch.utils.data import DataLoader
@@ -9,12 +10,18 @@ from torch.utils.data import DataLoader
 from accelerate import Accelerator
 from transformers import CLIPTextModel
 from diffusers import AutoencoderKL, UNet2DConditionModel, DPMSolverMultistepScheduler
-from diffusers.loaders import AttnProcsLayers
+from peft import LoraConfig, LoraModel, get_peft_model_state_dict
+from peft.tuners.lora import mark_only_lora_as_trainable
 
 from models.clip.tokenizer import MultiCLIPTokenizer
 from training.functional import TrainingStrategy, TrainingCallbacks, save_samples
 
 
+# https://github.com/huggingface/peft/blob/main/examples/lora_dreambooth/train_dreambooth.py
+UNET_TARGET_MODULES = ["to_q", "to_v", "query", "value"]
+TEXT_ENCODER_TARGET_MODULES = ["q_proj", "v_proj"]
+
+
 def lora_strategy_callbacks(
     accelerator: Accelerator,
     unet: UNet2DConditionModel,
@@ -27,7 +34,6 @@ def lora_strategy_callbacks(
     sample_output_dir: Path,
     checkpoint_output_dir: Path,
     seed: int,
-    lora_layers: AttnProcsLayers,
     max_grad_norm: float = 1.0,
     sample_batch_size: int = 1,
     sample_num_batches: int = 1,
@@ -57,7 +63,8 @@ def lora_strategy_callbacks(
     )
 
     def on_prepare():
-        lora_layers.requires_grad_(True)
+        mark_only_lora_as_trainable(unet.model, unet.peft_config.bias)
+        mark_only_lora_as_trainable(text_encoder.model, text_encoder.peft_config.bias)
 
     def on_accum_model():
         return unet
@@ -73,24 +80,44 @@ def lora_strategy_callbacks(
         yield
 
     def on_before_optimize(lr: float, epoch: int):
-        accelerator.clip_grad_norm_(lora_layers.parameters(), max_grad_norm)
+        accelerator.clip_grad_norm_(
+            itertools.chain(unet.parameters(), text_encoder.parameters()),
+            max_grad_norm
+        )
 
     @torch.no_grad()
     def on_checkpoint(step, postfix):
         print(f"Saving checkpoint for step {step}...")
 
         unet_ = accelerator.unwrap_model(unet, False)
-        unet_.save_attn_procs(
-            checkpoint_output_dir / f"{step}_{postfix}",
-            safe_serialization=True
+        text_encoder_ = accelerator.unwrap_model(text_encoder, False)
+
+        lora_config = {}
+        state_dict = get_peft_model_state_dict(unet, state_dict=accelerator.get_state_dict(unet))
+        lora_config["peft_config"] = unet.get_peft_config_as_dict(inference=True)
+
+        text_encoder_state_dict = get_peft_model_state_dict(
+            text_encoder, state_dict=accelerator.get_state_dict(text_encoder)
         )
+        text_encoder_state_dict = {f"text_encoder_{k}": v for k, v in text_encoder_state_dict.items()}
+        state_dict.update(text_encoder_state_dict)
+        lora_config["text_encoder_peft_config"] = text_encoder.get_peft_config_as_dict(inference=True)
+
+        accelerator.print(state_dict)
+        accelerator.save(state_dict, checkpoint_output_dir / f"{step}_{postfix}.pt")
+
         del unet_
+        del text_encoder_
 
     @torch.no_grad()
     def on_sample(step):
         unet_ = accelerator.unwrap_model(unet, False)
+        text_encoder_ = accelerator.unwrap_model(text_encoder, False)
+
         save_samples_(step=step, unet=unet_)
+
         del unet_
+        del text_encoder_
 
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -114,13 +141,34 @@ def lora_prepare(
     train_dataloader: DataLoader,
     val_dataloader: Optional[DataLoader],
     lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
-    lora_layers: AttnProcsLayers,
+    lora_rank: int = 4,
+    lora_alpha: int = 32,
+    lora_dropout: float = 0,
+    lora_bias: str = "none",
     **kwargs
 ):
-    lora_layers, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(
-        lora_layers, optimizer, train_dataloader, val_dataloader, lr_scheduler)
+    unet_config = LoraConfig(
+        r=lora_rank,
+        lora_alpha=lora_alpha,
+        target_modules=UNET_TARGET_MODULES,
+        lora_dropout=lora_dropout,
+        bias=lora_bias,
+    )
+    unet = LoraModel(unet_config, unet)
+
+    text_encoder_config = LoraConfig(
+        r=lora_rank,
+        lora_alpha=lora_alpha,
+        target_modules=TEXT_ENCODER_TARGET_MODULES,
+        lora_dropout=lora_dropout,
+        bias=lora_bias,
+    )
+    text_encoder = LoraModel(text_encoder_config, text_encoder)
+
+    text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(
+        text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler)
 
-    return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, {"lora_layers": lora_layers}
+    return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, {}
 
 
 lora_strategy = TrainingStrategy(