4 files changed, 32 insertions, 28 deletions

diff --git a/training/functional.py b/training/functional.py
index b830261..990c4cd 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -22,7 +22,6 @@ from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from models.clip.embeddings import ManagedCLIPTextEmbeddings, patch_managed_embeddings
 from models.clip.util import get_extended_embeddings
 from models.clip.tokenizer import MultiCLIPTokenizer
-from schedulers.scheduling_deis_multistep import DEISMultistepScheduler
 from training.util import AverageMeter
 
 
@@ -74,19 +73,12 @@ def make_grid(images, rows, cols):
     return grid
 
 
-def get_models(pretrained_model_name_or_path: str, noise_scheduler: str = "ddpm"):
-    if noise_scheduler == "deis":
-        noise_scheduler_cls = DEISMultistepScheduler
-    elif noise_scheduler == "ddpm":
-        noise_scheduler_cls = DDPMScheduler
-    else:
-        raise ValueError(f"noise_scheduler must be one of [\"ddpm\", \"deis\"], got {noise_scheduler}")
-
+def get_models(pretrained_model_name_or_path: str):
     tokenizer = MultiCLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer')
     text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder')
     vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder='vae')
     unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder='unet')
-    noise_scheduler = noise_scheduler_cls.from_pretrained(pretrained_model_name_or_path, subfolder='scheduler')
+    noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder='scheduler')
     sample_scheduler = UniPCMultistepScheduler.from_pretrained(
         pretrained_model_name_or_path, subfolder='scheduler')
 
@@ -232,9 +224,6 @@ def generate_class_images(
 
     del pipeline
 
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-
 
 def add_placeholder_tokens(
     tokenizer: MultiCLIPTokenizer,
@@ -274,26 +263,41 @@ def loss_step(
     latents = vae.encode(batch["pixel_values"]).latent_dist.sample()
     latents = latents * vae.config.scaling_factor
 
+    bsz = latents.shape[0]
+
     generator = torch.Generator(device=latents.device).manual_seed(seed + step) if eval else None
 
     # Sample noise that we'll add to the latents
-    noise = torch.randn(
-        latents.shape,
-        dtype=latents.dtype,
-        layout=latents.layout,
-        device=latents.device,
-        generator=generator
-    )
-    if low_freq_noise > 0:
-        noise += low_freq_noise * torch.randn(
+    if low_freq_noise == 0:
+        noise = torch.randn(
+            latents.shape,
+            dtype=latents.dtype,
+            layout=latents.layout,
+            device=latents.device,
+            generator=generator
+        )
+    else:
+        noise = (1 - low_freq_noise) * torch.randn(
+            latents.shape,
+            dtype=latents.dtype,
+            layout=latents.layout,
+            device=latents.device,
+            generator=generator
+        ) + low_freq_noise * torch.randn(
             latents.shape[0], latents.shape[1], 1, 1,
             dtype=latents.dtype,
             layout=latents.layout,
             device=latents.device,
             generator=generator
         )
+        # noise += low_freq_noise * torch.randn(
+        #     bsz, 1, 1, 1,
+        #     dtype=latents.dtype,
+        #     layout=latents.layout,
+        #     device=latents.device,
+        #     generator=generator
+        # )
 
-    bsz = latents.shape[0]
     # Sample a random timestep for each image
     timesteps = torch.randint(
         0,
diff --git a/training/optimization.py b/training/optimization.py
index 6c9a35d..7d8d55a 100644
--- a/training/optimization.py
+++ b/training/optimization.py
@@ -113,7 +113,7 @@ def get_scheduler(
 ):
     num_training_steps_per_epoch = math.ceil(
         num_training_steps_per_epoch / gradient_accumulation_steps
-    ) * gradient_accumulation_steps
+    )  # * gradient_accumulation_steps
     num_training_steps = train_epochs * num_training_steps_per_epoch
     num_warmup_steps = warmup_epochs * num_training_steps_per_epoch
 
diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py
index 0290327..e5e84c8 100644
--- a/training/strategy/dreambooth.py
+++ b/training/strategy/dreambooth.py
@@ -88,8 +88,8 @@ def dreambooth_strategy_callbacks(
 
     def on_prepare():
         unet.requires_grad_(True)
-        text_encoder.requires_grad_(True)
-        text_encoder.text_model.embeddings.requires_grad_(False)
+        text_encoder.text_model.encoder.requires_grad_(True)
+        text_encoder.text_model.final_layer_norm.requires_grad_(True)
 
         if ema_unet is not None:
             ema_unet.to(accelerator.device)
@@ -203,7 +203,7 @@ def dreambooth_prepare(
     lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
     **kwargs
 ):
-    return accelerator.prepare(text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler) + ({})
+    return accelerator.prepare(text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler) + ({},)
 
 
 dreambooth_strategy = TrainingStrategy(
diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 732cd74..bd0d178 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -130,7 +130,7 @@ def textual_inversion_strategy_callbacks(
             if lambda_ != 0:
                 w = text_encoder.text_model.embeddings.temp_token_embedding.weight
 
-                mask = torch.zeros(w.size(0), dtype=torch.bool)
+                mask = torch.zeros(w.shape[0], dtype=torch.bool)
                 mask[text_encoder.text_model.embeddings.temp_token_ids] = True
                 mask[zero_ids] = False