From 55586d20a76cc64cbb22331d86cdb7085fe093ea Mon Sep 17 00:00:00 2001 From: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com> Date: Thu, 12 Feb 2026 00:03:25 +0530 Subject: [PATCH] Fix operator precedence bug in x0 LVLB weights and NaN assertion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two bugs in DDPM.register_schedule(): 1. The x0 parameterization formula `2. * 1 - torch.Tensor(alphas_cumprod)` is evaluated as `(2. * 1) - torch.Tensor(alphas_cumprod)` due to operator precedence, yielding `2.0 - alphas_cumprod` instead of the intended `2.0 * (1 - alphas_cumprod)`. This produces incorrect LVLB weights that silently degrade training when using x0 parameterization. 2. The NaN guard `assert not isnan(...).all()` only fires when *every* element is NaN. A single NaN — which is enough to corrupt the loss — passes undetected. Changed to `.any()` so any NaN triggers the assertion. --- ldm/models/diffusion/ddpm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ldm/models/diffusion/ddpm.py b/ldm/models/diffusion/ddpm.py index bbedd04cf..03a665c78 100644 --- a/ldm/models/diffusion/ddpm.py +++ b/ldm/models/diffusion/ddpm.py @@ -160,13 +160,13 @@ def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps= lvlb_weights = self.betas ** 2 / ( 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)) elif self.parameterization == "x0": - lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod)) + lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * (1 - torch.Tensor(alphas_cumprod))) else: raise NotImplementedError("mu not supported") # TODO how to choose this term lvlb_weights[0] = lvlb_weights[1] self.register_buffer('lvlb_weights', lvlb_weights, persistent=False) - assert not torch.isnan(self.lvlb_weights).all() + assert not torch.isnan(self.lvlb_weights).any() @contextmanager def ema_scope(self, context=None):