fix qwen edit shape error

Updated citation format and added TMLR link.

.gitignore

@@ -98,3 +98,5 @@ cookiecutter-pypackage-env/
 # vscode settings
 .history/
 *.code-workspace
+.vscode/
+/.vscode

.vscode/settings.json

@@ -5,4 +5,8 @@
     "/PATH/TO/ComfyUI/",
     "/PATH/TO/ComfyUI/custom_nodes/"
   ],
+  "cursorpyright.analysis.extraPaths": [
+    "/PATH/TO/ComfyUI/",
+    "/PATH/TO/ComfyUI/custom_nodes/"
+  ],
 }

README.md

@@ -1,7 +1,7 @@
 <div align="center">
 
 # LanPaint: Universal Inpainting Sampler with "Think Mode"
-[![arXiv](https://img.shields.io/badge/Arxiv-2502.03491-b31b1b.svg?logo=arXiv)](https://arxiv.org/abs/2502.03491)
+[![TMLR PDF](https://img.shields.io/badge/TMLR-PDF-8A2BE2?logo=openreview&logoColor=white)](https://openreview.net/pdf?id=JPC8JyOUSW)
 [![Python Benchmark](https://img.shields.io/badge/🐍-Python_Benchmark-3776AB?logo=python)](https://github.com/scraed/LanPaintBench)
 [![ComfyUI Extension](https://img.shields.io/badge/ComfyUI-Extension-7B5DFF)](https://github.com/comfyanonymous/ComfyUI)
 [![Hugging Face](https://img.shields.io/badge/Hugging%20Face-yellow?logo=huggingface&logoColor=white)](https://huggingface.co/charrywhite/LanPaint)
@@ -12,7 +12,29 @@
 
 Universally applicable inpainting ability for every model. LanPaint sampler lets the model "think" through multiple iterations before denoising, enabling you to invest more computation time for superior inpainting quality.  
 
-This is the official implementation of ["Lanpaint: Training-Free Diffusion Inpainting with Exact and Fast Conditional Inference"](https://arxiv.org/abs/2502.03491), accepted by TMLR. The repository is for ComfyUI extension. Local Python benchmark code is published here: [LanPaintBench](https://github.com/scraed/LanPaintBench).
+This is the official implementation of ["LanPaint: Training-Free Diffusion Inpainting with Asymptotically Exact and Fast Conditional Sampling"](https://arxiv.org/abs/2502.03491), accepted by TMLR. The repository is for ComfyUI extension. Local Python benchmark code is published here: [LanPaintBench](https://github.com/scraed/LanPaintBench).
+
+## Citation
+
+```
+@article{
+zheng2025lanpaint,
+title={LanPaint: Training-Free Diffusion Inpainting with Asymptotically Exact and Fast Conditional Sampling},
+author={Candi Zheng and Yuan Lan and Yang Wang},
+journal={Transactions on Machine Learning Research},
+issn={2835-8856},
+year={2025},
+url={https://openreview.net/forum?id=JPC8JyOUSW},
+note={}
+}
+```
+
+**🎬 NEW: LanPaint now supports inpainting and outpainting based on Z-Image!**
+
+| Original | Masked | Inpainted |
+|:--------:|:------:|:---------:|
+| ![Original Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_21/Original_No_Mask.png) | ![Masked Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_21/Masked_Load_Me_in_Loader.png) | ![Inpainted Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_21/InPainted_Drag_Me_to_ComfyUI.png) |
+
 
 **🎬 NEW: LanPaint now supports video inpainting and outpainting based on Wan 2.2!**
 
@@ -36,9 +58,11 @@ Check our latest [Wan 2.2 Video Examples](#video-examples-beta), [Wan 2.2 Image
 - [How to Use Examples](#how-to-use-examples)
 - [Video Examples (Beta)](#video-examples-beta)
   - [Wan 2.2 Video Inpainting](#wan-22-video-inpainting)
+  - [Wan 2.2 5B Video Inpainting](#wan-22-5b-video-inpainting)
   - [Wan 2.2 Video Outpainting](#wan-22-video-outpainting)
   - [Resource Consumption](#resource-consumption)
 - [Image Examples](#image-examples)
+  - [Z-image](#example-z-image-inpaintlanpaint-k-sampler-5-steps-of-thinking)
   - [Hunyuan T2I](#example-hunyuan-t2i-inpaintlanpaint-k-sampler-5-steps-of-thinking)
   - [Wan 2.2 T2I](#example-wan22-inpaintlanpaint-k-sampler-5-steps-of-thinking)
   - [Wan 2.2 T2I with reference](#example-wan22-partial-inpaintlanpaint-k-sampler-5-steps-of-thinking)
@@ -53,7 +77,8 @@ Check our latest [Wan 2.2 Video Examples](#video-examples-beta), [Wan 2.2 Image
   - [Basic Sampler](#basic-sampler)
   - [Advanced Sampler](#lanpaint-ksampler-advanced)
   - [Tuning Guide](#lanpaint-ksampler-advanced-tuning-guide)
-- [Community Showcase](#community-showcase-) 
+- [Community Showcase](#community-showcase-)
+- [FAQ](#faq)
 - [Updates](#updates)
 - [ToDo](#todo)
 - [Citation](#citation)
@@ -114,6 +139,10 @@ LanPaint now supports video inpainting with Wan 2.2, enabling you to seamlessly
 
 You need to follow the ComfyUI version of [Wan2.2 T2V workflow](https://docs.comfy.org/tutorials/video/wan/wan2_2) to download and install the T2V model.
 
+### Wan 2.2 5B Video Inpainting 
+
+Similar to Wan 2.2 14B with slightly different workflow. [View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_17)
+
 ### Wan 2.2 Video Outpainting
 
 Extend your videos beyond their original boundaries with LanPaint's video outpainting capability based on Wan 2.2. This feature allows you to expand the canvas of your videos while maintaining coherent motion and context.
@@ -191,7 +220,7 @@ You need to follow the ComfyUI version of [Wan2.2 T2V workflow](https://docs.com
 **Model Used**: `wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors` and `wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors`.<br>
 **Processing Steps**:  20 sampling steps x 2 (LanPaint steps of thinking).</sub>
 
-**Note:** To further reduce VRAM requirements, we recommend loading CLIP on CPU.
+**Note:** Vram is required by the model, not LanPaint. To further reduce VRAM requirements, we recommend generating less frames and loading CLIP on CPU.
 
 ## Image Examples
 
@@ -212,6 +241,33 @@ We are excited to announce that LanPaint now supports Wan2.2 text to image gener
 
 You need to follow the ComfyUI version of [Wan2.2 T2V workflow](https://docs.comfy.org/tutorials/video/wan/wan2_2) to download and install the T2V model.
 
+### Example Z-image: InPaint(LanPaint K Sampler, 5 steps of thinking)
+LanPaint also supports inpainting with the Z-image text-to-image model.
+
+<details open>
+<summary>View Original / Masked / Inpainted Comparison</summary>
+
+| Original | Masked | Inpainted |
+|:--------:|:------:|:---------:|
+| ![Original Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_21/Original_No_Mask.png) | ![Masked Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_21/Masked_Load_Me_in_Loader.png) | ![Inpainted Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_21/InPainted_Drag_Me_to_ComfyUI.png) |
+
+</details>
+
+[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_21)
+
+<details open>
+<summary>View Z-image Outpainting (Original / Masked / Outpainted)</summary>
+
+| Original | Masked | Outpainted |
+|:--------:|:------:|:----------:|
+| ![Original Z-image Outpaint](https://github.com/scraed/LanPaint/blob/master/examples/Example_22/Original_No_Mask.png) | ![Masked Z-image Outpaint](https://github.com/scraed/LanPaint/blob/master/examples/Example_22/Masked_Load_Me_in_Loader.png) | ![Outpainted Z-image](https://github.com/scraed/LanPaint/blob/master/examples/Example_22/InPainted_Drag_Me_to_ComfyUI.png) |
+
+</details>
+
+[View Outpaint Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_22)
+
+You can download the Z-image model for ComfyUI from [Z-image](https://docs.comfy.org/zh-CN/tutorials/image/z-image/z-image-turbo).
+
 ### Example Wan2.2: Partial InPaint(LanPaint K Sampler, 5 steps of thinking)
 Sometimes we don't want to inpaint completely new content, but rather let the inpainted image reference the original image. One option to achieve this is to inpaint with an edit model like Qwen Image Edit. Another option is to perform a partial inpaint: allowing the diffusion process to start at some middle steps rather than from 0.
 
@@ -221,6 +277,7 @@ Sometimes we don't want to inpaint completely new content, but rather let the in
 
 You need to follow the ComfyUI version of [Wan2.2 T2V workflow](https://docs.comfy.org/tutorials/video/wan/wan2_2) to download and install the T2V model.
 
+
 ### Example Qwen Edit 2509: InPaint
 Check our latest updated [Mased Qwen Edit Workflow](https://github.com/scraed/LanPaint/tree/master/examples/Example_14) for Qwen Image Edit 2509. Download the model at [Qwen Image Edit 2509 Comfy](https://huggingface.co/Comfy-Org/Qwen-Image-Edit_ComfyUI/tree/main/split_files/diffusion_models).
 
@@ -379,6 +436,9 @@ Discover how the community is using LanPaint! Here are some user-created tutoria
 
 Submit a PR to add your tutorial/video here, or open an [Issue](https://github.com/scraed/LanPaint/issues) with details!
 
+## FAQ
+[Working togather with crop&stitch](https://github.com/scraed/LanPaint/issues/46)
+
 ## Updates
 - 2025/08/08
     - Add Qwen image support
@@ -407,17 +467,19 @@ Submit a PR to add your tutorial/video here, or open an [Issue](https://github.c
 - Try Implement Detailer
 - ~~Provide inference code on without GUI.~~ Check our local Python benchmark code [LanPaintBench](https://github.com/scraed/LanPaintBench).
 
+
 ## Citation
 
 ```
-@misc{zheng2025lanpainttrainingfreediffusioninpainting,
-      title={Lanpaint: Training-Free Diffusion Inpainting with Exact and Fast Conditional Inference}, 
-      author={Candi Zheng and Yuan Lan and Yang Wang},
-      year={2025},
-      eprint={2502.03491},
-      archivePrefix={arXiv},
-      primaryClass={eess.IV},
-      url={https://arxiv.org/abs/2502.03491}, 
+@article{
+zheng2025lanpaint,
+title={LanPaint: Training-Free Diffusion Inpainting with Asymptotically Exact and Fast Conditional Sampling},
+author={Candi Zheng and Yuan Lan and Yang Wang},
+journal={Transactions on Machine Learning Research},
+issn={2835-8856},
+year={2025},
+url={https://openreview.net/forum?id=JPC8JyOUSW},
+note={}
 }
 ```
 
@@ -425,4 +487,3 @@ Submit a PR to add your tutorial/video here, or open an [Issue](https://github.c
 
 
 
-

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "LanPaint"
-version = "1.4.2"
+version = "1.4.8"
 description = "Achieve seamless inpainting results without needing a specialized inpainting model."
 authors = [
   {name = "LanPaint", email = "czhengac@connect.ust.hk"}

src/LanPaint/lanpaint.py

@@ -13,6 +13,7 @@ class LanPaint():
         self.friction = Friction
         self.chara_beta = Beta
         self.img_dim_size = None
+
     def add_none_dims(self, array):
         # Create a tuple with ':' for the first dimension and 'None' repeated num_nones times
         index = (slice(None),) + (None,) * (self.img_dim_size-1)
@@ -42,6 +43,8 @@ class LanPaint():
             return self.inner_model.inner_model.model_sampling.noise_scaling(sigma.reshape([sigma.shape[0]] + [1] * (len(noise.shape) - 1)), noise, latent_image)
 
         x = x * (1 - latent_mask) +  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 * ( self.add_none_dims(abt)**0.5 + (1-self.add_none_dims(abt))**0.5 )

src/LanPaint/nodes.py

@@ -11,7 +11,8 @@ from comfy.samplers import *
 from comfy.model_base import ModelType
 from .utils import *
 from .lanpaint import LanPaint
-
+from comfy.model_base import WAN22
+import comfyui_version 
 
 def reshape_mask(input_mask, output_shape,video_inpainting=False):
     dims = len(output_shape) - 2
@@ -30,7 +31,9 @@ def reshape_mask(input_mask, output_shape,video_inpainting=False):
         
         # First reshape input_mask to have proper dimensions for video processing
         # Assume input is (frames, channels, height, width) -> (1, channels, frames, height, width)
-        input_mask = input_mask.permute(1, 0, 2, 3).unsqueeze(0)
+        ## if comfy version < 0.6.0
+        if comfyui_version.__version__ < "0.6.0":
+            input_mask = input_mask.permute(1, 0, 2, 3).unsqueeze(0)
         print('Video case - input_mask after reshaping:', input_mask.shape)
         # Ensure we have the correct 5D shape: (batch, channels, frames, height, width)
         batch_size, channels, frames, height, width = input_mask.shape
@@ -53,12 +56,15 @@ def reshape_mask(input_mask, output_shape,video_inpainting=False):
         # Handle batch dimension
         mask = repeat_to_batch_size(mask, output_shape[0])
     else:  # Original 2D image case
-        mask = torch.nn.functional.interpolate(input_mask, size=output_shape[-2:], mode=scale_mode)
+        if comfyui_version.__version__ < "0.6.0":
+            mask = torch.nn.functional.interpolate(input_mask, size=output_shape[-2:], mode=scale_mode)
+        else:
+            mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode)
         if mask.shape[1] < output_shape[1]:
             mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]]
         mask = repeat_to_batch_size(mask, output_shape[0])
     
-    print('resize mask',mask.shape,type(mask),torch.max(mask),torch.min(mask))
+
     return mask
 def prepare_mask(noise_mask, shape, device,video_inpainting=False):
     return reshape_mask(noise_mask, shape,video_inpainting).to(device)
@@ -85,6 +91,10 @@ class CFGGuider_LanPaint:
         self.inner_model, self.conds, self.loaded_models = comfy.sampler_helpers.prepare_sampling(self.model_patcher, noise.shape, self.conds, self.model_options)
         device = self.model_patcher.load_device
 
+        if isinstance(self.inner_model, WAN22):
+            print("WAN22 detected")
+            self.inner_model.extra_conds = super(WAN22, self.inner_model).extra_conds
+
         if denoise_mask is not None:
             video_inpainting = self.model_options.get("video_inpainting", False)
             denoise_mask = prepare_mask(denoise_mask, noise.shape, device, video_inpainting)
@@ -125,8 +135,9 @@ class KSamplerX0Inpaint:
 
         IS_FLUX = self.inner_model.inner_model.model_type == ModelType.FLUX
         IS_FLOW = self.inner_model.inner_model.model_type == ModelType.FLOW
-        print("model type", self.inner_model.inner_model.model_type, "IS_FLUX", IS_FLUX, "IS_FLOW", IS_FLOW)
-        print("sigma", torch.mean(sigma).item(), torch.min(sigma).item(), torch.max(sigma).item())
+        #print("model class", type(self.inner_model.inner_model))
+        #print("model type", self.inner_model.inner_model.model_type, "IS_FLUX", IS_FLUX, "IS_FLOW", IS_FLOW)
+        #print("sigma", torch.mean(sigma).item(), torch.min(sigma).item(), torch.max(sigma).item())
         # unify the notations into variance exploding diffusion model
         if IS_FLUX or IS_FLOW:
             Flow_t = sigma
@@ -404,8 +415,16 @@ class MaskBlend:
         # for each mask pixel, find out the nearest 1 pixel, and set the mask value to the distance between the two pixels
         # check the size of mask and image1, image2, if not the same, assert error
         if image1.shape[1] != image2.shape[1] or image1.shape[2] != image2.shape[2]:
-            raise ValueError("Make sure your image size is a multiple of 8. Otherwise the mask will not be aligned with the output image.")
-        
+            raise ValueError(
+                "Image size mismatch: Image1 and Image2 must have the same dimensions.\n"
+                "Additionally, ensure both images have width and height that are multiples of 8.\n"
+                "This is required because VAE decode always generates images with dimensions that are multiples of 8.\n"
+                "If your input images are not multiples of 8, a size mismatch will occur during the decoding process.\n"
+                "Please resize your images using an image resize node to ensure compatibility.\n"
+                "Current sizes - Image1: {}x{}, Image2: {}x{}".format(
+                    image1.shape[2], image1.shape[1], image2.shape[2], image2.shape[1]
+                )
+            )
         mask = mask.float()
         mask = torch.nn.functional.max_pool2d(mask, kernel_size=blend_overlap, stride=1, padding=blend_overlap//2)
         # apply Gaussian blur with kernel size blend_overlap