Update

1 files changed, 27 insertions, 23 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,