From d2b43185173e4fd5dd5d6705f6975e75d475f915 Mon Sep 17 00:00:00 2001 From: Nithin GK <24860019+Nithin-GK@users.noreply.github.com> Date: Wed, 28 Jun 2023 11:17:56 -0700 Subject: [PATCH] Update models.py Fixed bug with classifier-free guidance sampling, the number of input channels of LDMs is 4. --- models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/models.py b/models.py index c90eeba7..8a38fc0d 100644 --- a/models.py +++ b/models.py @@ -252,6 +252,7 @@ def forward_with_cfg(self, x, t, y, cfg_scale): Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance. """ # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb + C=x.shape[1] half = x[: len(x) // 2] combined = torch.cat([half, half], dim=0) model_out = self.forward(combined, t, y) @@ -259,7 +260,7 @@ def forward_with_cfg(self, x, t, y, cfg_scale): # three channels by default. The standard approach to cfg applies it to all channels. # This can be done by uncommenting the following line and commenting-out the line following that. # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:] - eps, rest = model_out[:, :3], model_out[:, 3:] + eps, rest = model_out[:, :C], model_out[:, C:] cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0) half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps) eps = torch.cat([half_eps, half_eps], dim=0)