Newer Model support for LanPaint on diffusers

[chaserhkj]

Hi, I wish to use newer models (ZImage, Qwen-Image, etc.) with LanPaint on diffusers. The main reason of using diffusers is mostly batch processing, as well as exploring some other automated workflows like externally invoked through MCP by LLM agents or auto detailer. ComfyUI's principle design makes its API and automated integration really a pain to work with.

I know the existence of LanPaintBench, but that project does not support newer models and is not really maintained.

I have started a new WIP proper port of lanpaint onto diffusers, and I will be willing to maintain it onwards as well. I am opening this issue mainly to have a tracked reference here, but I also have a few questions regarding the implementation and algorithms:

1. What are the biggest change in the underlying math for the newer models to diverge from the old math that works for SD?
2. It seems that some maths of LanPaintBench and ComfyUI versions are done slightly different, I am not 100% sure if they are equivalent, does the ComfyUI version of the algorithm have more validity since it is more maintained?
3. The sigma functions sigma_x and sigma_y seem to be different from LanPaintBench version and ComfyUI version. Is this just a change of heuristics? Does this mean the heuristics used in the ComfyUI version is better?

Thanks a lot for this amazing project!

[scraed]

Hi, you are welcome to maintain a diffusers version of LanPaint!

1: There are no difference in math for different models. The only thing you need to take are of is different notation: VE, VP, and Flow model, which I have maintained three corresponding time variable sigma, abt, and flow_t in Comfy version.

    if IS_FLUX or IS_FLOW:
            Flow_t = sigma #Comfy's sigma is flow time here
            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())


        else:
            VE_Sigma = sigma #Comfy use VE notation for non flow models
            abt = 1/( 1+VE_Sigma**2 )
            Flow_t = (1-abt)**0.5 / ( (1-abt)**0.5 + abt**0.5  )

The ComfyUI version also convert VE, VP, and Flow into VP notation first, then do Langevin dynamics, finally convert back.

if IS_FLUX or IS_FLOW:
            x_t = x * ( self.add_none_dims(abt)**0.5 + (1-self.add_none_dims(abt))**0.5 )
        else:
            x_t = x / ( 1+self.add_none_dims(VE_Sigma)**2 )**0.5 # switch to variance perserving x_t values

The LanPaintBench only handles VE notation. I would suggest use VP notation as it have better compatibility with Flow models.

2: The difference between Comfy and Bench is mainly on notation, one works in VP notation, one works in VE. This leads to different Coef_C and D. I would suggest VP notation for compatibility with flow modes.

3: In Bench, m is usually 1 and alpha is recommend to be 0. So both sigma xy are just (1-abt). They will be multiplied with stepsize as schedule

    def sigma_y(self, abt):
        return (1-abt)**self.m
    def sigma_x(self, abt):
        return (1 - abt + abt * self.alpha)**self.m

In Comfy, the schedule is done explicitly by

step_size = self.step_size * (1 - abt)

Therefore the sigmas are just constants.

def sigma_x(self, abt):
        # the time scale for the x_t update
        return abt**0
    def sigma_y(self, abt):
        beta = self.chara_beta * abt ** 0
        return beta

These sigmas are actually legacy code I haven't clean up. Forgive me they are not consistent.

4. Difference between VP (Comfy) and VE (Bench) in Langevin Coefficients:

For bench's VE notation

        D_x = (2 * (1 + sigma**2) )**0.5
        D_y = (2 * (1 + sigma**2) )**0.5
...
        def Coef_C(x_t):
            x0 = self.x0_evalutation(x_t, score, sigma, abt, args)
            C = ( x0 - x_t ) /  (1-abt) + A * x_t
            return C

For Comfy's VP notation

        D_x = (2 * abt**0 )**0.5
        D_y = (2 * abt**0 )**0.5
...
        def Coef_C(x_t):
            x0 = x_t + score(x_t)
            C = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t
            return C, x0

Finally, be careful when you compute score(x_t). The following code is from Comfy's implementation. The VP notation need to be switched back to the model's desired notation before feeding into the model, which is implemented by the if else. self.inner_model in the following code predicts clean image x_0. If your model predicts eps or velocity, you need to switch to x_0 prediction with some math. ( you can ignore x_0_BIG, which just aims to allow a different CFG scale for score_y. Treating x_0_BIG = x_0 is totally fine. )

    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.
            x = x_t / ( abt**0.5 + (1-abt)**0.5 ) # switch to Gaussian flow matching
            x_0, x_0_BIG = self.inner_model(x, self.remove_none_dims(tflow), model_options=model_options, seed=seed)
        else:
            x = x_t * ( 1+sigma**2 )**0.5 # switch to variance exploding
            x_0, x_0_BIG = self.inner_model(x, self.remove_none_dims(sigma), model_options=model_options, seed=seed)

        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

[alvarocardenasc]

I’m very interested in the Python version as well.

[kobechenyang]

Tried use Claude code to generate the code from the comfyui workflow.
Feels close but not quite there. I am using the 4B flux2 klein model though. will try again with 9B one.

import math
import torch
import time
import os
import requests
import numpy as np
from io import BytesIO
from PIL import Image

# CUDA optimizations
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = True

from diffusers import Flux2KleinPipeline

# Import LanPaint
try:
    from LanPaint.lanpaint import LanPaint as LanPaintEngine
except ImportError:
    print("LanPaint not installed. Run:")
    print("  pip install git+https://github.com/godnight10061/LanPaint.git@docs/pure-python-usage")
    exit(1)


# Test image URLs from LanPaint repo
INPUT_IMAGE_URL = "https://raw.githubusercontent.com/scraed/LanPaint/master/examples/Example_24/Original_No_Mask.png"
MASK_IMAGE_URL = "https://raw.githubusercontent.com/scraed/LanPaint/master/examples/Example_24/Masked_Load_Me_in_Loader.png"

# Workflow parameters (from Flux2_Klein_inpainting.json)
TEST_PROMPT = "change building's window light color to blue"
NEGATIVE_PROMPT = ""
NUM_STEPS = 20
CFG_SCALE = 5.0
SEED = 0

# LanPaint parameters (from workflow)
LANPAINT_NSTEPS = 2
LANPAINT_LAMBDA = 8.0
LANPAINT_STEPSIZE = 0.2
LANPAINT_BETA = 1.0
LANPAINT_FRICTION = 15.0
LANPAINT_EARLY_STOP = 1
BLEND_OVERLAP = 9


# ---------------------------------------------------------------------------
# Flux2 Scheduler (from ComfyUI comfy_extras/nodes_flux.py)
# ---------------------------------------------------------------------------
def generalized_time_snr_shift(t, mu, sigma):
    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)


def compute_empirical_mu(image_seq_len, num_steps):
    a1, b1 = 8.73809524e-05, 1.89833333
    a2, b2 = 0.00016927, 0.45666666
    if image_seq_len > 4300:
        return float(a2 * image_seq_len + b2)
    m_200 = a2 * image_seq_len + b2
    m_10 = a1 * image_seq_len + b1
    a = (m_200 - m_10) / 190.0
    b = m_200 - 200.0 * a
    return float(a * num_steps + b)


def get_flux2_sigmas(num_steps, width, height):
    """Compute Flux2 sigma schedule (matching ComfyUI Flux2Scheduler)."""
    seq_len = round(width * height / (16 * 16))
    mu = compute_empirical_mu(seq_len, num_steps)
    timesteps = torch.linspace(1, 0, num_steps + 1)
    sigmas = torch.tensor([generalized_time_snr_shift(t.item(), mu, 1.0) if t > 0 else 0.0 for t in timesteps])
    return sigmas


# ---------------------------------------------------------------------------
# Time variable helpers (from LanPaint timevars.py / ComfyUI nodes.py)
# ---------------------------------------------------------------------------
def make_current_times_flow(flow_t):
    """Compute (ve_sigma, abt, flow_t) from flow time."""
    one_minus_t = 1.0 - flow_t
    abt = one_minus_t ** 2 / (one_minus_t ** 2 + flow_t ** 2)
    ve_sigma = flow_t / torch.clamp(one_minus_t, min=1e-9)
    return ve_sigma, abt, flow_t


# ---------------------------------------------------------------------------
# Model Wrapper — matches ComfyUI's KSamplerX0Inpaint + CFGGuider_LanPaint
# ---------------------------------------------------------------------------
class ModelSampling:
    """Matches ComfyUI's CONST model_sampling for flow matching."""
    def noise_scaling(self, sigma, noise, latent_image):
        sigma = sigma.reshape([sigma.shape[0]] + [1] * (len(noise.shape) - 1))
        return sigma * noise + (1.0 - sigma) * latent_image


class FluxTransformerWrapper:
    """
    Wraps the diffusers Flux2Klein transformer to match LanPaint's expected interface.

    LanPaint expects: model(x, sigma, model_options=None, seed=None) -> (x0, x0_BIG)

    This wrapper handles:
    - Reference latent concatenation (Flux2 img2img mechanism)
    - CFG with dual scales (cfg for x0, cfg_BIG=1.0 for x0_BIG)
    - Velocity → x0 conversion
    - Latent packing/unpacking
    """

    def __init__(
        self,
        pipeline,
        pos_prompt_embeds: torch.Tensor,
        pos_text_ids: torch.Tensor,
        neg_prompt_embeds: torch.Tensor,
        neg_text_ids: torch.Tensor,
        ref_latent: torch.Tensor,
        cfg_scale: float = 5.0,
        cfg_big: float = 1.0,
    ):
        self.transformer = pipeline.transformer
        self.pipeline = pipeline

        self.pos_prompt_embeds = pos_prompt_embeds
        self.pos_text_ids = pos_text_ids
        self.neg_prompt_embeds = neg_prompt_embeds
        self.neg_text_ids = neg_text_ids
        self.cfg_scale = cfg_scale
        self.cfg_big = cfg_big

        # Pre-compute packed reference latent and its IDs
        self.ref_packed = self.pipeline._pack_latents(ref_latent)  # (B, seq, C)
        self.ref_ids = self.pipeline._prepare_image_ids([ref_latent])  # (1, seq, 4)

        # Store model dtype (LanPaint may promote inputs to float32 via mask ops)
        self.model_dtype = next(pipeline.transformer.parameters()).dtype

        # LanPaint compatibility
        self.inner_model = self
        self.model_sampling = ModelSampling()

    def __call__(self, x, sigma, model_options=None, seed=None):
        """
        Call transformer with CFG and return (x0, x0_BIG).

        Args:
            x: Noisy latent [B, C, H, W] (patchified, BN-normalized)
            sigma: Flow time tensor (scalar or [B])

        Returns:
            (x0, x0_BIG): Two x0 predictions with different CFG scales
        """
        batch_size = x.shape[0]

        # Cast x to model dtype (LanPaint may have promoted to float32 via mask ops)
        x = x.to(dtype=self.model_dtype)

        # Prepare timestep (sigma is flow_t in [0, 1])
        timestep = sigma.clone()
        if timestep.dim() == 0:
            timestep = timestep.unsqueeze(0)
        timestep = timestep.expand(batch_size).to(x.device, dtype=self.model_dtype)

        # Pack noise latent: (B, C, H, W) -> (B, H*W, C)
        noise_packed = self.pipeline._pack_latents(x)
        noise_seq_len = noise_packed.shape[1]

        # Create noise latent IDs (created on CPU by default, move to device)
        noise_ids = self.pipeline._prepare_latent_ids(x).to(device=x.device, dtype=x.dtype)

        # Expand pre-computed reference latent + IDs to batch and move to device
        ref_packed = self.ref_packed.to(x.device, dtype=x.dtype).expand(batch_size, -1, -1)
        ref_ids = self.ref_ids.to(x.device, dtype=x.dtype).expand(batch_size, -1, -1)

        # Concatenate noise + reference along sequence dimension
        hidden_states = torch.cat([noise_packed, ref_packed], dim=1)
        img_ids = torch.cat([noise_ids, ref_ids], dim=1)

        # Run model for both cond and uncond
        with torch.no_grad():
            # Positive (conditional) prediction
            cond_out = self.transformer(
                hidden_states=hidden_states,
                timestep=timestep,
                encoder_hidden_states=self.pos_prompt_embeds.to(x.device, dtype=x.dtype).expand(batch_size, -1, -1),
                txt_ids=self.pos_text_ids.to(x.device, dtype=x.dtype).expand(batch_size, -1, -1),
                img_ids=img_ids,
                guidance=None,
                return_dict=False,
            )[0]

            # Negative (unconditional) prediction
            uncond_out = self.transformer(
                hidden_states=hidden_states,
                timestep=timestep,
                encoder_hidden_states=self.neg_prompt_embeds.to(x.device, dtype=x.dtype).expand(batch_size, -1, -1),
                txt_ids=self.neg_text_ids.to(x.device, dtype=x.dtype).expand(batch_size, -1, -1),
                img_ids=img_ids,
                guidance=None,
                return_dict=False,
            )[0]

        # Slice to keep only noise tokens (discard reference tokens)
        cond_noise = cond_out[:, :noise_seq_len]
        uncond_noise = uncond_out[:, :noise_seq_len]

        # Unpack to spatial (B, seq, C) -> (B, C, H, W) using position IDs
        cond_v = self.pipeline._unpack_latents_with_ids(cond_noise, noise_ids)
        uncond_v = self.pipeline._unpack_latents_with_ids(uncond_noise, noise_ids)

        # CFG: velocity = uncond + scale * (cond - uncond)
        v_cfg = uncond_v + self.cfg_scale * (cond_v - uncond_v)
        v_cfg_big = uncond_v + self.cfg_big * (cond_v - uncond_v)

        # Convert velocity to x0: x0 = x_t - sigma * v
        t = sigma.view(-1, 1, 1, 1).to(x.device, dtype=x.dtype)
        x0 = x - t * v_cfg
        x0_big = x - t * v_cfg_big

        return x0, x0_big


# ---------------------------------------------------------------------------
# VAE encode / decode
# ---------------------------------------------------------------------------
def _patchify_latents(latents):
    """(B, C, H, W) -> (B, C*4, H//2, W//2)"""
    B, C, H, W = latents.shape
    return latents.view(B, C, H // 2, 2, W // 2, 2).permute(0, 1, 3, 5, 2, 4).reshape(B, C * 4, H // 2, W // 2)


def _unpatchify_latents(latents):
    """(B, C*4, H//2, W//2) -> (B, C, H, W)"""
    B, C, H, W = latents.shape
    return latents.reshape(B, C // 4, 2, 2, H, W).permute(0, 1, 4, 2, 5, 3).reshape(B, C // 4, H * 2, W * 2)


def encode_image(image, vae, device, dtype):
    """Encode PIL image to Flux2 latent space (patchified + BN normalized)."""
    import torchvision.transforms as T
    img_tensor = T.ToTensor()(image).unsqueeze(0).to(device, dtype=dtype)
    img_tensor = 2.0 * img_tensor - 1.0

    with torch.no_grad():
        latent = vae.encode(img_tensor).latent_dist.sample()
        latent = _patchify_latents(latent)
        bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latent.device, latent.dtype)
        bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(latent.device, latent.dtype)
        latent = (latent - bn_mean) / bn_std
    return latent


def decode_latent(latent, vae):
    """Decode Flux2 latent to PIL image."""
    with torch.no_grad():
        bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latent.device, latent.dtype)
        bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps).to(latent.device, latent.dtype)
        latent = latent * bn_std + bn_mean
        latent = _unpatchify_latents(latent)
        img = vae.decode(latent, return_dict=False)[0]
        img = ((img + 1.0) / 2.0).clamp(0, 1)
    img_np = img[0].permute(1, 2, 0).cpu().float().numpy()
    return Image.fromarray((img_np * 255).astype("uint8"))


# ---------------------------------------------------------------------------
# Image / mask helpers
# ---------------------------------------------------------------------------
def download_image(url):
    print(f"Downloading: {url}")
    response = requests.get(url)
    response.raise_for_status()
    return Image.open(BytesIO(response.content))


def get_flux_dimensions(width, height):
    """Scale to ~1MP with dimensions divisible by 16."""
    max_area = 1024 ** 2
    multiple_of = 16
    aspect = width / height
    new_w = round((max_area * aspect) ** 0.5)
    new_h = round((max_area / aspect) ** 0.5)
    return new_w // multiple_of * multiple_of, new_h // multiple_of * multiple_of


def prepare_denoise_mask(mask_image, latent_shape, device, dtype):
    """
    Prepare denoise mask from alpha channel of loaded image.

    ComfyUI convention: denoise_mask 1 = inpaint, 0 = keep
    The mask image has: black (0) = inpaint region, white (1) = keep region
    So we invert: denoise_mask = 1 - mask_value
    """
    import torchvision.transforms as T
    mask_tensor = T.ToTensor()(mask_image).to(device, dtype=dtype)  # [1, H, W]
    # Invert: black (0) → denoise_mask 1 (inpaint), white (1) → denoise_mask 0 (keep)
    denoise_mask = 1.0 - mask_tensor

    _, _, latent_h, latent_w = latent_shape
    denoise_mask = torch.nn.functional.interpolate(
        denoise_mask.unsqueeze(0), size=(latent_h, latent_w), mode="nearest"
    )  # [1, 1, H, W]
    return denoise_mask


def prepare_mask_for_image(mask_image, image_size):
    """Prepare mask at image resolution for post-processing blend."""
    import torchvision.transforms as T
    mask_tensor = T.ToTensor()(mask_image)  # [1, H, W]
    # Invert: black (0) → 1 (inpaint area), white (1) → 0 (keep area)
    return 1.0 - mask_tensor  # [1, H, W]


# ---------------------------------------------------------------------------
# Mask blend post-processing (from ComfyUI LanPaint MaskBlend node)
# ---------------------------------------------------------------------------
def gaussian_kernel(kernel_size):
    sigma = (kernel_size - 1) / 4
    x = torch.arange(kernel_size).float() - kernel_size // 2
    y = torch.arange(kernel_size).float() - kernel_size // 2
    x_grid, y_grid = torch.meshgrid(x, y, indexing='ij')
    kernel = torch.exp(-(x_grid ** 2 + y_grid ** 2) / (2 * sigma ** 2))
    return kernel / kernel.sum()


def mask_blend(original_image, inpainted_image, mask, blend_overlap=9):
    """
    Blend original and inpainted images using dilated + blurred mask.

    Args:
        original_image: PIL Image (original, before inpainting)
        inpainted_image: PIL Image (output from denoising)
        mask: Tensor [1, H, W] where 1 = inpaint area, 0 = keep area
        blend_overlap: Kernel size for dilation and blur (must be odd)

    Returns:
        PIL Image (blended result)
    """
    import torchvision.transforms as T

    img1 = T.ToTensor()(original_image).unsqueeze(0)  # [1, 3, H, W]
    img2 = T.ToTensor()(inpainted_image).unsqueeze(0)

    # Resize img2 to match img1 if needed
    if img1.shape != img2.shape:
        img2 = torch.nn.functional.interpolate(img2, size=img1.shape[2:], mode="bilinear", align_corners=False)

    mask = mask.float().unsqueeze(0)  # [1, 1, H, W]

    # Dilate mask with max_pool2d
    mask = torch.nn.functional.max_pool2d(mask, kernel_size=blend_overlap, stride=1, padding=blend_overlap // 2)

    # Blur with Gaussian kernel
    kernel = gaussian_kernel(blend_overlap)
    kernel = kernel[None, None, ...]  # [1, 1, K, K]
    mask = torch.nn.functional.conv2d(mask, kernel, padding=blend_overlap // 2)

    # Blend
    blended = img1 * (1 - mask) + img2 * mask
    blended = blended.clamp(0, 1)

    img_np = blended[0].permute(1, 2, 0).numpy()
    return Image.fromarray((img_np * 255).astype("uint8"))


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16

    print("=" * 60)
    print("LanPaint + Flux2-Klein Inpainting Test")
    print("=" * 60)

    # [1/6] Download test images
    print("\n[1/6] Downloading test images...")
    input_image = download_image(INPUT_IMAGE_URL).convert("RGB")
    mask_image = download_image(MASK_IMAGE_URL)

    # Extract mask from alpha channel (ComfyUI LoadImage output 1)
    mask_rgba = mask_image.convert("RGBA")
    _, _, _, alpha = mask_rgba.split()
    a_array = np.array(alpha)
    if a_array.min() < 250:
        mask_bw = alpha  # Use alpha channel as mask
    else:
        mask_bw = mask_image.convert("L")

    # Scale to 1MP (matching ComfyUI ImageScaleToTotalPixels)
    width, height = get_flux_dimensions(input_image.size[0], input_image.size[1])
    print(f"Original size: {input_image.size} -> Flux size: {width}x{height}")

    input_image = input_image.resize((width, height), Image.Resampling.LANCZOS)
    mask_bw = mask_bw.resize((width, height), Image.Resampling.NEAREST)

    os.makedirs("outputs", exist_ok=True)
    input_image.save("outputs/lanpaint_test_input.png")
    mask_bw.save("outputs/lanpaint_test_mask.png")
    print("Saved input and mask to outputs/")

    # [2/6] Load pipeline
    print("\n[2/6] Loading Flux2-Klein pipeline...")
    pipeline = Flux2KleinPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-klein-4B",
        torch_dtype=dtype,
    ).to(device)

    vae = pipeline.vae
    print("Pipeline loaded successfully")

    # [3/6] Encode prompts (positive + negative)
    print("\n[3/6] Encoding prompts...")
    print(f"Positive: {TEST_PROMPT}")
    print(f"Negative: '{NEGATIVE_PROMPT}'")

    pos_embeds, pos_text_ids = pipeline.encode_prompt(
        prompt=TEST_PROMPT, device=device, num_images_per_prompt=1,
    )
    neg_embeds, neg_text_ids = pipeline.encode_prompt(
        prompt=NEGATIVE_PROMPT, device=device, num_images_per_prompt=1,
    )

    # [4/6] Encode image + prepare mask
    print("\n[4/6] Encoding image and preparing mask...")
    input_latent = encode_image(input_image, vae, device, dtype)
    print(f"Latent shape: {input_latent.shape}")

    # Denoise mask: 1 = inpaint, 0 = keep (ComfyUI convention)
    denoise_mask = prepare_denoise_mask(mask_bw, input_latent.shape, device, dtype)
    print(f"Denoise mask shape: {denoise_mask.shape}")

    # [5/6] Build model wrapper and run LanPaint
    print("\n[5/6] Running LanPaint inpainting...")
    print(f"Steps: {NUM_STEPS}, CFG: {CFG_SCALE}, LanPaint NSteps: {LANPAINT_NSTEPS}")

    # Create model wrapper with reference latent + CFG
    model_wrapper = FluxTransformerWrapper(
        pipeline=pipeline,
        pos_prompt_embeds=pos_embeds,
        pos_text_ids=pos_text_ids,
        neg_prompt_embeds=neg_embeds,
        neg_text_ids=neg_text_ids,
        ref_latent=input_latent,
        cfg_scale=CFG_SCALE,
        cfg_big=1.0,  # For FLUX models, cfg_BIG is always 1.0
    )

    # Create LanPaint engine (IS_FLUX=True, IS_FLOW=False — matching ComfyUI ModelType.FLUX)
    lanpaint_engine = LanPaintEngine(
        Model=model_wrapper,
        NSteps=LANPAINT_NSTEPS,
        Friction=LANPAINT_FRICTION,
        Lambda=LANPAINT_LAMBDA,
        Beta=LANPAINT_BETA,
        StepSize=LANPAINT_STEPSIZE,
        IS_FLUX=True,
        IS_FLOW=False,
    )

    # Compute Flux2 sigma schedule
    sigmas = get_flux2_sigmas(NUM_STEPS, width, height).to(device, dtype=dtype)
    print(f"Sigmas: {sigmas[0].item():.4f} -> {sigmas[-1].item():.4f} ({len(sigmas)} values)")

    # Initialize noise
    torch.manual_seed(SEED)
    noise = torch.randn_like(input_latent)

    # Initial noisy latent: x = sigma * noise + (1 - sigma) * latent_image
    x = model_wrapper.model_sampling.noise_scaling(sigmas[0:1], noise, input_latent)

    start_time = time.time()

    # Euler sampling loop with LanPaint (matching KSAMPLER.sample in nodes.py)
    total_steps = len(sigmas) - 1
    for i in range(total_steps):
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        sigma_tensor = sigma.unsqueeze(0)

        print(f"  Step {i + 1}/{total_steps} (sigma={sigma.item():.4f})", end="\r")

        # Compute time variables for LanPaint
        current_times = make_current_times_flow(flow_t=sigma_tensor)

        # Binarize denoise mask and compute latent_mask (1=keep, 0=inpaint)
        denoise_mask_bin = (denoise_mask > 0.5).float()
        latent_mask = 1.0 - denoise_mask_bin

        # LanPaint step (with early stop for last N steps)
        if total_steps - i <= LANPAINT_EARLY_STOP:
            out = lanpaint_engine(x, input_latent, noise, sigma_tensor, latent_mask,
                                  current_times, {}, SEED, n_steps=0)
        else:
            out = lanpaint_engine(x, input_latent, noise, sigma_tensor, latent_mask,
                                  current_times, {}, SEED)

        # Euler step: x_next = x + (sigma_next - sigma) * velocity
        # velocity = (x - x0) / sigma (from CONST: x0 = x - sigma * v)
        dt = sigma_next - sigma
        d = (x - out) / sigma.clamp(min=1e-9)
        x = x + d * dt

    generation_time = time.time() - start_time
    print(f"\nGeneration completed in {generation_time:.2f} seconds")

    # [6/6] Decode and blend
    print("\n[6/6] Decoding and blending...")
    output_image = decode_latent(x.to(dtype=dtype), vae)

    # Mask blend post-processing (matching LanPaint_MaskBlend node)
    image_mask = prepare_mask_for_image(mask_bw, input_image.size)
    final_image = mask_blend(input_image, output_image, image_mask, blend_overlap=BLEND_OVERLAP)

    # Save outputs
    output_image.save("outputs/lanpaint_test_raw_output.png")
    final_image.save("outputs/lanpaint_test_output.png")
    print(f"Raw output saved to: outputs/lanpaint_test_raw_output.png")
    print(f"Blended output saved to: outputs/lanpaint_test_output.png")

    print("\n" + "=" * 60)
    print("Test completed!")
    print("=" * 60)


if __name__ == "__main__":
    main()

[alvarocardenasc]

@kobechenyang Did you pass the JSON to Claude Code, or how did you manage to generate that?

[chaserhkj]

I do have a partial implementation in the linked repo but I've been busy on a few other projects lately so might need to come back to this later.

@kobechenyang I would not place too much hope on the Claude generated code. The math notations are all over the place and it would be nearly impossible to verify if the math it did is correct there.

AI agents are notoriously bad at generating code that leverages math heavily, and they rarely get the math correct. What's worse is that they often produce code that still gives you somewhat correctly-formed outputs, but the underlying logic and values are completely broken. They are often much better if you are just generating code that changes APIs or moves data around. If you really want to try this I would suggest to do some prompt engineering by providing utils.py to the agent and explicitly ask it to use functions from that file to do math. All math functions in that file is framework-independent and I plan to just copy it to my implementation anyways. This way you can let the agent do the API connecting that they are often much better at and avoid messing with the core math.

[kobechenyang]

@alvarocardenasc I was using the exported api json.
@chaserhkj I do use the Lanpaint as the base project so it should have the full context. Unfortunate I don't know too much about the math part to verify its correctness.
Today I tried with the 9B base model. it actually generate pretty good result. I will do more testing.

[kobechenyang]

another example

[alvarocardenasc]

@kobechenyang looks good, could you share the code?

[scraed]

Looks good. the code is importing the most math heavy part directly from LanPaint.lanpaint import LanPaint as LanPaintEngine. It might differ a bit but not much.

[scraed]

@kobechenyang looks good, could you share the code?

The code has already been pasted here. Putting it together with LanPaint's original code should work. #77 (comment)

[kobechenyang]

Yes the code has been shared above. @scraed is it correct that it LanPaint would only work with base model not on any distilled model? for example, it won't work here with the 4B flux2 klein model?

[scraed]

Yes the code has been shared above. @scraed is it correct that it LanPaint would only work with base model not on any distilled model? for example, it won't work here with the 4B flux2 klein model?

No. That is not correct. The performance might not that good on distilled model, depending on their quality of distillation. But it works.