diff --git a/examples/Example_7/InPainted_Drag_Me_to_ComfyUI.png b/examples/Example_7/InPainted_Drag_Me_to_ComfyUI.png index 38bbecc..6455f9c 100644 Binary files a/examples/Example_7/InPainted_Drag_Me_to_ComfyUI.png and b/examples/Example_7/InPainted_Drag_Me_to_ComfyUI.png differ diff --git a/examples/Example_7/Masked_Load_Me_in_Loader.png b/examples/Example_7/Masked_Load_Me_in_Loader.png index 582acbd..717e80f 100644 Binary files a/examples/Example_7/Masked_Load_Me_in_Loader.png and b/examples/Example_7/Masked_Load_Me_in_Loader.png differ diff --git a/examples/Example_8/InPainted_Drag_Me_to_ComfyUI.png b/examples/Example_8/InPainted_Drag_Me_to_ComfyUI.png index efe6ce4..d57190c 100644 Binary files a/examples/Example_8/InPainted_Drag_Me_to_ComfyUI.png and b/examples/Example_8/InPainted_Drag_Me_to_ComfyUI.png differ diff --git a/examples/Example_9/InPainted_Drag_Me_to_ComfyUI.png b/examples/Example_9/InPainted_Drag_Me_to_ComfyUI.png new file mode 100644 index 0000000..723a021 Binary files /dev/null and b/examples/Example_9/InPainted_Drag_Me_to_ComfyUI.png differ diff --git a/examples/Example_9/Masked_Load_Me_in_Loader.png b/examples/Example_9/Masked_Load_Me_in_Loader.png new file mode 100644 index 0000000..9fb1fb7 Binary files /dev/null and b/examples/Example_9/Masked_Load_Me_in_Loader.png differ diff --git a/examples/Example_9/Original_No_Mask.png b/examples/Example_9/Original_No_Mask.png new file mode 100644 index 0000000..f0a435a Binary files /dev/null and b/examples/Example_9/Original_No_Mask.png differ diff --git a/examples/InpaintChara_10.jpg b/examples/InpaintChara_10.jpg index 2256fa6..a1ee9d9 100644 Binary files a/examples/InpaintChara_10.jpg and b/examples/InpaintChara_10.jpg differ diff --git a/examples/InpaintChara_11.jpg b/examples/InpaintChara_11.jpg index 7783f55..64ed4c2 100644 Binary files a/examples/InpaintChara_11.jpg and b/examples/InpaintChara_11.jpg differ diff --git a/examples/InpaintChara_12.jpg b/examples/InpaintChara_12.jpg index 2649252..f7996ad 100644 Binary files a/examples/InpaintChara_12.jpg and b/examples/InpaintChara_12.jpg differ diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py index b3ad8b3..882f759 100644 --- a/src/LanPaint/lanpaint.py +++ b/src/LanPaint/lanpaint.py @@ -18,15 +18,12 @@ class LanPaint(): #self.LanPaint_Cap_Sigma = CapSigma self.chara_beta = Beta - def __call__(self, x, latent_image, noise, sigma, Sigmas, latent_mask, current_times, model_options, seed): - self.VE_Sigmas = Sigmas + def __call__(self, x, latent_image, noise, sigma, latent_mask, current_times, model_options, seed): self.latent_image = latent_image self.noise = noise return self.LanPaint(x, sigma, latent_mask, current_times, model_options, seed, self.IS_FLUX, self.IS_FLOW) def LanPaint(self, x, sigma, latent_mask, current_times, model_options, seed, IS_FLUX, IS_FLOW): - VE_Sigma, abt = current_times - sigma_ind = torch.argmin(torch.abs(self.VE_Sigmas - torch.mean( VE_Sigma ))) - VE_Sigma_next = self.VE_Sigmas[ sigma_ind + 1 ] * VE_Sigma**0 + VE_Sigma, abt, Flow_t = current_times step_size = self.step_size * (1 - abt) @@ -36,42 +33,40 @@ class LanPaint(): x = x * (1 - latent_mask) + self.inner_model.inner_model.scale_latent_inpaint(x=x, sigma=sigma, noise=self.noise, latent_image=self.latent_image)* latent_mask if IS_FLUX or IS_FLOW: - x_t = x * ( 1 + VE_Sigma[:, None,None,None]) + x_t = x * ( abt[:, None,None,None]**0.5 + (1-abt[:, None,None,None])**0.5 ) else: - x_t = x #/ ( 1+sigma**2 )**0.5 # switch to variance perserving x_t values + x_t = x / ( 1+VE_Sigma[:, None,None,None]**2 )**0.5 # switch to variance perserving x_t values ############ LanPaint Iterations Start ############### # after noise_scaling, noise = latent_image + noise * sigma, which is x_t in the variance exploding diffusion model notation for the known region. args = None for i in range(self.n_steps): - score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = abt[:, None,None,None], sigma = VE_Sigma[:, None,None,None], model_options = model_options, seed = seed ) + score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = abt[:, None,None,None], sigma = VE_Sigma[:, None,None,None], tflow = Flow_t[:, None,None,None], model_options = model_options, seed = seed ) x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size , current_times, sigma_x = self.sigma_x(abt)[:, None,None,None], sigma_y = self.sigma_y(abt)[:, None,None,None], args = args) if IS_FLUX or IS_FLOW: - x = x_t / ( 1 + VE_Sigma[:, None,None,None] ) + x = x_t / ( abt[:, None,None,None]**0.5 + (1-abt[:, None,None,None])**0.5 ) else: - x = x_t #/ ( 1+sigma**2 )**0.5 # switch to variance perserving x_t values + x = x_t * ( 1+VE_Sigma[:, None,None,None]**2 )**0.5 # switch to variance perserving x_t values ############ LanPaint Iterations End ############### # out is x_0 out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed) out = out * (1-latent_mask) + self.latent_image * latent_mask return out - def score_model(self, x_t, y, mask, abt, sigma, model_options, seed): + def score_model(self, x_t, y, mask, abt, sigma, tflow, model_options, seed): lamb = self.chara_lamb if self.IS_FLUX or self.IS_FLOW: # compute t for flow model, with a small epsilon compensating for numerical error. - t_flow = sigma[:, 0,0,0] / ( 1 + sigma[:, 0,0,0] - 5e-3 * sigma[:, 0,0,0] ) - x_0, x_0_BIG = self.inner_model(x_t / ( 1 + sigma ), t_flow, model_options=model_options, seed=seed) + x = x_t / ( abt**0.5 + (1-abt)**0.5 ) # switch to Gaussian flow matching + x_0, x_0_BIG = self.inner_model(x, tflow[:, 0,0,0], model_options=model_options, seed=seed) else: - x_0, x_0_BIG = self.inner_model(x_t, sigma[:, 0,0,0], model_options=model_options, seed=seed) + x = x_t * ( 1+sigma**2 )**0.5 # switch to variance exploding + x_0, x_0_BIG = self.inner_model(x, sigma[:, 0,0,0], model_options=model_options, seed=seed) - e_t = x_t / ((1 - abt) ** 0.5 * (1 + sigma**2) ** 0.5 )- (abt ** 0.5 / (1 - abt) ** 0.5) * x_0 - e_t_BIG = x_t / ((1 - abt) ** 0.5 * (1 + sigma**2) ** 0.5 )- (abt ** 0.5 / (1 - abt) ** 0.5) * x_0_BIG - - score_x = -e_t - score_y = - (1 + lamb) * ( x_t/ ((1 + sigma**2) ** 0.5 *(1 - abt)**0.5) - abt**0.5 /(1 - abt)**0.5 * y ) + lamb * e_t_BIG + score_x = -(x_t - x_0) + score_y = - (1 + lamb) * ( x_t - y ) + lamb * (x_t - x_0_BIG) return score_x * (1 - mask) + score_y * mask def sigma_x(self, abt): # the time scale for the x_t update @@ -89,10 +84,10 @@ class LanPaint(): if torch.mean(dtx) <= 0.: return x_t, args # ------------------------------------------------------------------------- - # A: Update epsilon (score estimate and noise initialization) + # Compute the Langevin dynamics update in variance perserving notation # ------------------------------------------------------------------------- x0 = self.x0_evalutation(x_t, score, sigma, args) - C = x0 / (1-abt) + C = abt**0.5 * x0 / (1-abt) A = A_x * (1-mask) + A_y * mask D = D_x * (1-mask) + D_y * mask dt = dtx * (1-mask) + dty * mask @@ -115,7 +110,7 @@ class LanPaint(): def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y): # ------------------------------------------------------------------------- # Unpack current times parameters (sigma and abt) - sigma, abt = current_times + sigma, abt, flow_t = current_times sigma = sigma[:, None,None,None] abt = abt[:, None,None,None] # Compute time step (dtx, dty) for x and y branches. @@ -141,16 +136,14 @@ class LanPaint(): Gamma_x = Gamma_hat_x / (dtx/2) Gamma_y = Gamma_hat_y / (dty/2) - D_x = (2 * (1 + sigma**2) )**0.5 - D_y = (2 * (1 + sigma**2) )**0.5 - + #D_x = (2 * (1 + sigma**2) )**0.5 + #D_y = (2 * (1 + sigma**2) )**0.5 + D_x = (2 * abt**0 )**0.5 + D_y = (2 * abt**0 )**0.5 return sigma, abt, dtx/2, dty/2, Gamma_x, Gamma_y, A_x, A_y, D_x, D_y def x0_evalutation(self, x_t, score, sigma, args): - score_model = score(x_t) - eps_model = -score_model - - x0 = x_t - sigma * eps_model + x0 = x_t + score(x_t) return x0 \ No newline at end of file diff --git a/src/LanPaint/nodes.py b/src/LanPaint/nodes.py index 9473576..9518cc3 100644 --- a/src/LanPaint/nodes.py +++ b/src/LanPaint/nodes.py @@ -91,12 +91,16 @@ class KSamplerX0Inpaint: # unify the notations into variance exploding diffusion model if IS_FLUX or IS_FLOW: + Flow_t = sigma + abt = (1 - Flow_t)**2 / ((1 - Flow_t)**2 + Flow_t**2 ) + VE_Sigma = Flow_t / (1 - Flow_t) + #print("t", torch.mean( sigma ).item(), "VE_Sigma", torch.mean( VE_Sigma ).item()) - VE_Sigma = sigma / ( 1 - sigma + 5e-3 * sigma) - self.VE_Sigmas = self.sigmas / ( 1 - self.sigmas + 5e-3 * self.sigmas ) + else: VE_Sigma = sigma - self.VE_Sigmas = self.sigmas + abt = 1/( 1+VE_Sigma**2 ) + Flow_t = (1-abt)**0.5 / ( (1-abt)**0.5 + abt**0.5 ) if denoise_mask is not None: if "denoise_mask_function" in model_options: @@ -106,11 +110,11 @@ class KSamplerX0Inpaint: latent_mask = 1 - denoise_mask - abt = 1/( 1+VE_Sigma**2 ) + - current_times = (VE_Sigma, abt) + current_times = (VE_Sigma, abt, Flow_t) - out = self.PaintMethod(x, self.latent_image, self.noise, sigma, self.VE_Sigmas, latent_mask, current_times, model_options, seed) + out = self.PaintMethod(x, self.latent_image, self.noise, sigma, latent_mask, current_times, model_options, seed) else: out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed) @@ -248,9 +252,9 @@ class LanPaint_KSampler(): model.LanPaint_StepSize = 0.15 model.LanPaint_Lambda = 8.0 - model.LanPaint_Beta = 1. + model.LanPaint_Beta = 1.0 model.LanPaint_NumSteps = LanPaint_NumSteps - model.LanPaint_Friction = 15 + model.LanPaint_Friction = 15. if LanPaint_PromptMode == "Image First": model.LanPaint_cfg_BIG = cfg else: