From 2a67dd353f5eacf092d85bbd0a17e31f1d75fd7a Mon Sep 17 00:00:00 2001 From: scraed Date: Fri, 2 May 2025 21:14:54 +0800 Subject: [PATCH] update epxm1Dx --- src/LanPaint/nodes.py | 48 +++++++++++++++++++++++++++++++++++++++---- 1 file changed, 44 insertions(+), 4 deletions(-) diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py index 0c695bf..343139e 100644 --- a/src/LanPaint/nodes.py +++ b/src/LanPaint/nodes.py @@ -309,8 +309,16 @@ class KSamplerX0Inpaint: eps_denoise = eps_model_x * (1 - mask) + eps_model_y * mask return eps_denoise def eps_with_momentum(self, eps, eps_model, Gamma_hat_x, Gamma_hat_y, mask): - zeta_x = 2.0 / (Gamma_hat_x + 1e-4) * (1 - torch.exp(-0.5 * Gamma_hat_x)) - zeta_y = 2.0 / (Gamma_hat_y + 1e-4) * (1 - torch.exp(-0.5 * Gamma_hat_y)) + def epxm1Dx(x): + # Compute the (exp(x) - 1) / x term with a small value to avoid division by zero. + result = torch.special.expm1(x) / x + # replace NaN or inf values with 0 + result = torch.where(torch.isfinite(result), result, torch.zeros_like(result)) + mask = torch.abs(x) < 1e-2 + result = torch.where(mask, 1 + x/2. + x**2 / 6., result) + return result + zeta_x = epxm1Dx(-Gamma_hat_x / 2.) + zeta_y = epxm1Dx(-Gamma_hat_y / 2.) eps_bar_x = self.alpha * ( zeta_x * (eps - eps_model) ) + eps_model eps_bar_y = self.alpha * ( zeta_y * (eps - eps_model) ) + eps_model # Form the denoised epsilon using self.alpha (assumed to be 1/Ψ) @@ -334,17 +342,49 @@ class KSamplerX0Inpaint: Z_y = torch.exp(-0.5 * Gamma_hat_y) * Z + (1 - torch.exp(-Gamma_hat_y)) ** 0.5 * Z_q Z_next = Z_x * (1 - mask) + Z_y * mask + def epxm1Dx(x): + # Compute the (exp(x) - 1) / x term with a small value to avoid division by zero. + result = torch.special.expm1(x) / x + # replace NaN or inf values with 0 + result = torch.where(torch.isfinite(result), result, torch.zeros_like(result)) + mask = torch.abs(x) < 1e-2 + result = torch.where(mask, 1 + x/2. + x**2 / 6., result) + return result + # Compute the combined noise update following the scheme: - Z_comb_x = ( + zeta_x = epxm1Dx(-Gamma_hat_x / 2.) + zeta_y = epxm1Dx(-Gamma_hat_y / 2.) + def noise_combination(Gamma_hat, Z, Z_q, Z_q_avg): + exp_term = torch.special.expm1(-Gamma_hat / 2.) + tanh_term = torch.tanh(Gamma_hat / 4.) + + term1 = - exp_term / 2**0.5 * Z + term2 = - exp_term * tanh_term ** 0.5 / 2**0.5 * Z_q + term3 = torch.sqrt((Gamma_hat / 2.) - 2. * tanh_term) * Z_q_avg + + numerator = term1 + term2 + term3 + denominator = torch.sqrt(exp_term + (Gamma_hat / 2.)) + return numerator / denominator + # Compute the combined noise update following the scheme: + Z_comb_x = (1-zeta_x)**0.5 * noise_combination(Gamma_hat_x, Z, Z_q, Z_q_avg) + Z_comb_y = (1-zeta_y)**0.5 * noise_combination(Gamma_hat_y, Z, Z_q, Z_q_avg) + + + Z_comb_x_old = ( (1 - torch.exp(-Gamma_hat_x / 2)) / torch.sqrt(Gamma_hat_x + 1e-4) * (Z + torch.sqrt(torch.tanh(Gamma_hat_x / 4)) * Z_q) + torch.sqrt(1 - (4 / (Gamma_hat_x + 1e-4)) * torch.tanh(Gamma_hat_x / 4)) * Z_q_avg ) - Z_comb_y = ( + Z_comb_y_old = ( (1 - torch.exp(-Gamma_hat_y / 2)) / torch.sqrt(Gamma_hat_y + 1e-4) * (Z + torch.sqrt(torch.tanh(Gamma_hat_y / 4)) * Z_q) + torch.sqrt(1 - (4 / (Gamma_hat_y + 1e-4)) * torch.tanh(Gamma_hat_y / 4)) * Z_q_avg ) + + print("The same Zx", torch.allclose(Z_comb_x, Z_comb_x_old, rtol=1e-4, atol=1e-4)) + print("The same Zy", torch.allclose(Z_comb_y, Z_comb_y_old, rtol=1e-4, atol=1e-4)) + + Z_comb = Z_comb_x * (1 - mask) + Z_comb_y * mask # Combine with an additional noise term using self.alpha.