Update AltDiffusion

examples/AltDiffusion/dreambooth.py

@@ -11,6 +11,7 @@
 from pathlib import Path
 
 import torch
+from torch.cuda.amp import autocast as autocast
 from torch.utils.data import Dataset
 from PIL import Image
 from torchvision import transforms
@@ -34,7 +35,7 @@
 train_text_encoder = False
 train_only_unet = True
 
-num_train_epochs = 500
+num_train_epochs = 10
 batch_size = 4
 learning_rate = 5e-6
 adam_beta1 = 0.9
@@ -197,20 +198,23 @@ def collate_fn(examples):
         if with_prior_preservation:
             x, x_prior = torch.chunk(x, 2, dim=0)
             c, c_prior = torch.chunk(c, 2, dim=0)
-            loss, _ = model(x, c)
+            with autocast():
+                loss, _ = model(x, c)
             prior_loss, _ = model(x_prior, c_prior)
             loss = loss + prior_loss_weight * prior_loss
         else:
-            loss, _ = model(x, c)
+            with autocast():
+                loss, _ = model(x, c)
 
         print('*'*20, "loss=", str(loss.detach().item()))
 
-        loss.backward()
-        optimizer.step()
-        optimizer.zero_grad()
+        with autocast():
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
 
 ## mkdir ./checkpoints/DreamBooth and copy ./checkpoints/AltDiffusion to ./checkpoints/DreamBooth/AltDiffusion
 ## overwrite model.ckpt for latter usage
-chekpoint_path = './checkpoints/DreamBooth/AltDiffusion/model.ckpt'
+chekpoint_path = './checkpoints/AltDiffusion/dreambooth_model.ckpt'
 torch.save(model.state_dict(), chekpoint_path)
 

flagai/model/mm/AltDiffusion.py

@@ -1319,6 +1319,7 @@ def p_losses(self, x_start, cond, t, noise=None):
 
         loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
 
+        t = t.cpu()
         logvar_t = self.logvar[t].to(self.device)
         loss = loss_simple / torch.exp(logvar_t) + logvar_t
         # loss = loss_simple / torch.exp(self.logvar) + self.logvar
@@ -1932,4 +1933,4 @@ def normal_kl(mean1, logvar1, mean2, logvar2):
     ]
 
     return 0.5 * (-1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) +
-                  ((mean1 - mean2)**2) * torch.exp(-logvar2))
+                  ((mean1 - mean2)**2) * torch.exp(-logvar2))

requirements.txt

@@ -3,3 +3,4 @@ PyYAML==5.4.1
 deepspeed==0.6.5
 flash-attn==1.0.2
 bminf
+torch