Merge pull request #1 from visual-gen/fix

Fix

Author: DTennant

Diff for: semanticist/engine/diffusion_trainer.py

@@ -34,6 +34,7 @@ def __init__(
         warmup_epochs=100,
         warmup_steps=None,
         warmup_lr_init=0,
+        decay_steps=None,
         batch_size=32,
         eval_bs=32,
         test_bs=64,
@@ -130,7 +131,7 @@ def __init__(
             if self.accelerator.is_main_process:
                 print(f"Effective batch size is {effective_bs}")
 
-            self.g_optim = create_optimizer(self.model, weight_decay=0.05, learning_rate=lr, accelerator=self.accelerator)
+            self.g_optim = create_optimizer(self.model, weight_decay=0.05, learning_rate=lr,) # accelerator=self.accelerator)
 
             if warmup_epochs is not None:
                 warmup_steps = warmup_epochs * len(self.train_dl)
@@ -142,6 +143,7 @@ def __init__(
                 lr_min,
                 warmup_steps,
                 warmup_lr_init,
+                decay_steps,
                 cosine_lr
             )
             self.accelerator.register_for_checkpointing(self.g_sched)
@@ -232,7 +234,7 @@ def _load_checkpoint(self, ckpt_path=None):
             print(f"Loaded checkpoint from {ckpt_path}")
 
     def train(self, config=None):
-        n_parameters = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        n_parameters = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
         if self.accelerator.is_main_process:
             print(f"number of learnable parameters: {n_parameters//1e6}M")
         if config is not None:
@@ -293,7 +295,6 @@ def train(self, config=None):
                     self.accelerator.backward(loss)
                     if self.accelerator.sync_gradients and self.max_grad_norm is not None:
                         self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
-                    self.accelerator.unwrap_model(self.model).cancel_gradients_encoder(epoch)
                     self.g_optim.step()
                     if self.g_sched is not None:
                         self.g_sched.step_update(self.steps)
@@ -355,7 +356,7 @@ def evaluate(self):
                         img = batch
 
                     with self.accelerator.autocast():
-                        rec = self.model(img, targets, sample=True, inference_with_n_slots=self.test_num_slots, cfg=1.0)
+                        rec = self.model(img, sample=True, inference_with_n_slots=self.test_num_slots, cfg=1.0)
                     imgs_and_recs = torch.stack((img.to(rec.device), rec), dim=0)
                     imgs_and_recs = rearrange(imgs_and_recs, "r b ... -> (b r) ...")
                     imgs_and_recs = imgs_and_recs.detach().cpu().float()
@@ -373,7 +374,7 @@ def evaluate(self):
 
                     if self.cfg != 1.0:
                         with self.accelerator.autocast():
-                            rec = self.model(img, targets, sample=True, inference_with_n_slots=self.test_num_slots, cfg=self.cfg)
+                            rec = self.model(img, sample=True, inference_with_n_slots=self.test_num_slots, cfg=self.cfg)
 
                         imgs_and_recs = torch.stack((img.to(rec.device), rec), dim=0)
                         imgs_and_recs = rearrange(imgs_and_recs, "r b ... -> (b r) ...")
@@ -417,7 +418,7 @@ def process_batch(cfg_value, save_dir, header):
                         targets = targets.to(self.device, non_blocking=True)
 
                     with self.accelerator.autocast():
-                        recs = self.model(imgs, targets, sample=True, inference_with_n_slots=self.test_num_slots, cfg=cfg_value)
+                        recs = self.model(imgs, sample=True, inference_with_n_slots=self.test_num_slots, cfg=cfg_value)
 
                     psnr_val = psnr(recs, imgs, data_range=1.0)
                     ssim_val = ssim(recs, imgs, data_range=1.0)

Diff for: semanticist/engine/gpt_trainer.py

@@ -39,6 +39,7 @@ def __init__(
         warmup_epochs=100,
         warmup_steps=None,
         warmup_lr_init=0,
+        decay_steps=None,
         batch_size=32,
         cache_bs=8,
         test_bs=100,
@@ -137,6 +138,7 @@ def __init__(
                 lr_min,
                 warmup_steps,
                 warmup_lr_init,
+                decay_steps,
                 cosine_lr
             )
             self.accelerator.register_for_checkpointing(self.g_sched)

Diff for: semanticist/stage1/diffuse_slot.py

@@ -28,6 +28,7 @@ def __init__(
         torch.nn.init.normal_(self.null_cond, std=.02)
         self.autoenc_cond_embedder = nn.Linear(autoenc_dim, self.hidden_size)
         self.y_embedder = nn.Identity()
+        self.cond_drop_prob = 0.1
 
         self.use_repa = use_repa
         self._repa_hook = None
@@ -39,7 +40,21 @@ def embed_cond(self, autoenc_cond, drop_mask=None):
         # autoenc_cond: (N, K, D)
         # drop_ids: (N)
         # self.null_cond: (1, K, D)
-        autoenc_cond_drop = torch.where(drop_mask[:, :, None], autoenc_cond, self.null_cond)
+        batch_size = autoenc_cond.shape[0]
+        if drop_mask is None:
+            # randomly drop all conditions, for classifier-free guidance
+            if self.training:
+                drop_ids = (
+                    torch.rand(batch_size, 1, 1, device=autoenc_cond.device)
+                    < self.cond_drop_prob
+                )
+                autoenc_cond_drop = torch.where(drop_ids, self.null_cond, autoenc_cond)
+            else:
+                autoenc_cond_drop = autoenc_cond
+        else:
+            # randomly drop some conditions according to the drop_mask (N, K)
+            # True means keep
+            autoenc_cond_drop = torch.where(drop_mask[:, :, None], autoenc_cond, self.null_cond)
         return self.autoenc_cond_embedder(autoenc_cond_drop)
 
     def forward(self, x, t, autoenc_cond, drop_mask=None):
@@ -75,7 +90,7 @@ def forward_with_cfg(self, x, t, autoenc_cond, drop_mask, y=None, cfg_scale=1.0)
         # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
         half = x[: len(x) // 2]
         combined = torch.cat([half, half], dim=0)
-        model_out = self.forward(combined, t, autoenc_cond, drop_mask, y)
+        model_out = self.forward(combined, t, autoenc_cond, drop_mask)
         eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
         cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
         half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)

Diff for: train_net.py

@@ -1,7 +1,7 @@
 import os.path as osp
 import argparse
 from omegaconf import OmegaConf
-from semanticist.engine.util import instantiate_from_config
+from semanticist.engine.trainer_utils import instantiate_from_config
 from semanticist.utils.device_utils import configure_compute_backend
 
 def train():