Compare commits

...
129 Commits
Author SHA1 Message Date
scraed 75e606579c bump version 2026-03-26 17:10:41 +08:00
scraed 6440a7f433 unpack 2 bug fix for custom advanced node 2026-03-26 17:09:35 +08:00
scraed b3dc7e7551 bump version 2026-03-23 19:27:52 +08:00
scraed 2522687dd5 Klein workflow update 2026-03-23 19:26:45 +08:00
scraed f3ef9dedf9 update Klein workflow for resizing problem 2026-03-23 19:24:39 +08:00
scraed f53442c8d5 Revise README for v1.5.0 release notes
Updated README to include details about v1.5.0 bug fix and new features.
2026-03-03 11:46:56 +08:00
scraed 24051fb151 update README 2026-03-02 22:28:29 +08:00
scraed a76a5da0c5 bump version 2026-03-02 22:25:37 +08:00
scraed 21a6b155e8 fix the langevin output x has not been used issue 2026-03-02 14:16:32 +08:00
scraed 8bc20d5e00 Update README with additional resources and structure
Reorganized README content for clarity and added links to Diffusers
2026-02-16 01:03:30 +08:00
scraed 84b63adb11 replace flux example mask for bug fixing 2026-02-07 12:15:29 +08:00
scraed 725fb2ba37 Replace Klein example image 2026-02-07 11:56:03 +08:00
scraed 4f19674299 Update README.md 2026-02-04 14:40:08 +08:00
scraed 558084416d version 1.4.13 2026-02-02 23:49:35 +08:00
scraed a984dcb118 Fix flux klein example workflow 2026-02-02 23:48:33 +08:00
scraed 3f9193997a Update README.md 2026-02-02 20:36:55 +08:00
scraed 7372428114 Update README.md 2026-02-02 20:18:57 +08:00
scraed ac6eece747 Update Discord link in README.md 2026-02-02 20:15:45 +08:00
scraed 95b65ed1b1 Merge pull request #74 from godnight10061/fix/semantic-stop-tail
Add semantic inner-loop early stop mechanism
2026-01-30 11:44:51 +08:00
scraed 30cfd2f2e4 add contribution info 2026-01-30 11:40:02 +08:00
scraed 2e3a566bb4 Keep a minimal set of test 2026-01-30 11:23:51 +08:00
scraed 8fe1a9d312 Move parameters to the last for backward compatibility 2026-01-30 10:45:40 +08:00
scraed 4fcbfc7b5b Merge branch 'master' into pr/74 2026-01-30 10:20:59 +08:00
scraed 651ddc2225 fix typo 2026-01-30 00:48:22 +08:00
scraed bd4606f27a bump version to 1.4.12 2026-01-30 00:46:18 +08:00
scraed 8a72294f27 Z Image Base support 2026-01-30 00:45:20 +08:00
godnight10061 9ddd5f613e Map legacy min_steps into patience 2026-01-26 21:40:36 +08:00
godnight10061 39c90bbf62 Revert benchmark workflow 2026-01-26 21:16:05 +08:00
godnight10061 6ad4a18537 Refactor earlystop to reduce duplication 2026-01-26 18:09:21 +08:00
godnight10061 5e7d472738 Use LangevinState in semantic stop test mocks 2026-01-26 14:39:47 +08:00
godnight10061 d73f3db265 Document distance_fn hook 2026-01-26 14:02:07 +08:00
godnight10061 cd06069b52 Avoid early-stop when no inpaint region 2026-01-26 10:50:39 +08:00
godnight10061 114d091965 Simplify LanPaint inner-loop glue 2026-01-26 02:50:36 +08:00
godnight10061 06da285a09 Fix LangevinState backward compat 2026-01-25 23:42:04 +08:00
godnight10061 8f52a1d378 Use NamedTuple for Langevin state 2026-01-25 20:32:51 +08:00
godnight10061 c7962bf784 Refactor semantic early-stop metric 2026-01-25 10:57:29 +08:00
godnight10061 a895d2a758 Tighten semantic early-stop robustness 2026-01-24 12:43:23 +08:00
godnight10061 34ded92d8d Simplify semantic early-stop knobs 2026-01-24 10:19:37 +08:00
scraed 1b87239553 Move Coef_C definition above advance_time
Reordered the Coef_C function definition to appear before advance_time for improved code organization and readability.
2026-01-24 10:19:37 +08:00
scraed cbc9f3abfd Refactor early-stop logic into LanPaintEarlyStopper class
Moved the early-stop logic from lanpaint.py into a new src/LanPaint/earlystop.py module as the LanPaintEarlyStopper class. This improves code organization and maintainability by encapsulating early-stop behavior, reducing complexity in the main LanPaint class.
2026-01-24 10:19:37 +08:00
godnight10061 7d46147887 Guard torch.nn.functional import 2026-01-24 10:19:37 +08:00
godnight10061 c5fe6351ea Harden semantic-stop distance hook 2026-01-24 10:19:37 +08:00
godnight10061 ac7a5a1f2f Fix semantic-stop distance_fn dispatch 2026-01-24 10:19:37 +08:00
godnight10061 73ab537e86 Make CI-only imports lighter 2026-01-24 10:19:37 +08:00
godnight10061 6fd5f60383 Adjust SHO regression test for x0 return 2026-01-24 10:19:36 +08:00
godnight10061 9c9b5555fc Tighten semantic stop in mid-noise steps 2026-01-24 10:19:36 +08:00
godnight10061 450a25e5e8 Refine semantic inner stop 2026-01-24 10:19:36 +08:00
godnight10061 e7080f9574 Test ring changes prevent premature freeze 2026-01-24 10:19:36 +08:00
godnight10061 2086f00603 Skip score model after semantic convergence
# Conflicts:
#	src/LanPaint/lanpaint.py
2026-01-24 10:19:36 +08:00
godnight10061 896729ccd4 Fix semantic early-stop distance 2026-01-24 10:19:36 +08:00
godnight10061 207860edc3 Avoid semantic early-stop at low sigma 2026-01-24 10:19:36 +08:00
godnight10061 4074473b90 Add per-case quality gates to benchmark 2026-01-24 10:19:36 +08:00
godnight10061 11c1490140 Make benchmark workflow test semantic early stop 2026-01-24 10:19:36 +08:00
godnight10061 7c79712fc8 Make LPIPS benchmark long-running e2e 2026-01-24 10:19:36 +08:00
godnight10061 fb35a1d032 Add ImageNet LPIPS benchmark workflow 2026-01-24 10:19:36 +08:00
godnight10061 060fcc4475 Add semantic early-stop for Langevin iterations 2026-01-24 10:19:36 +08:00
scraed aedb908a3f Add version tuple check for ComfyUI version handling
Introduced a helper function to compare ComfyUI versions as tuples instead of strings, improving reliability of version-dependent logic in reshape_mask. Updated all version checks to use the new tuple-based comparison.
2026-01-22 16:22:41 +08:00
scraed f9718ea3e1 Update Flux 2 klein image example link
Corrected the anchor link for the 'Flux 2 klein' image example to reference '2 steps of thinking' instead of '5 steps of thinking'.
2026-01-18 00:55:57 +08:00
scraed 0bafe2117a Merge branch 'master' of https://github.com/scraed/LanPaint 2026-01-18 00:46:59 +08:00
scraed 2c1a23777d Update example steps in README
Changed the number of steps in the Flux 2 klein InPaint example from 5 to 2 to reflect the correct workflow.
2026-01-18 00:46:56 +08:00
scraed dde4c82463 Bump version from 1.4.10 to 1.4.11 2026-01-18 00:46:14 +08:00
scraed 72cb484c40 Add Flux 2 Klein inpainting example and workflow
Added a new image inpainting example for Flux 2 Klein, including workflow JSON, sample images (original, masked, inpainted), and updated README with documentation and links. This provides users with a reference workflow and visual results for the Flux 2 Klein model.
2026-01-18 00:39:22 +08:00
scraed 0d15d15cf3 Update Discord link in README.md 2026-01-13 23:34:02 +08:00
scraed 1393a46d67 Bump version from 1.4.9 to 1.4.10 2026-01-13 22:17:42 +08:00
scraed 27421363bf Merge pull request #71 from godnight10061/fix/issue69-nan-multivariatenormal
Avoid MultivariateNormal crash on non-finite dynamics
2026-01-13 22:16:34 +08:00
scraed 01400a541d Update test_sho_regression.py 2026-01-13 22:13:30 +08:00
scraed c981354387 Update test_sho_regression.py 2026-01-13 21:46:56 +08:00
scraed c92c8bfb1f Update test_sho_regression.py 2026-01-13 21:38:49 +08:00
scraed 9334025801 Update test_sho_regression.py 2026-01-13 21:31:15 +08:00
scraed 144893ecd2 fall back to overdamped update if nan appears 2026-01-13 18:41:27 +08:00
scraed bce8505ce2 Revert "Avoid MultivariateNormal crash on non-finite dynamics"
This reverts commit 9432a34c38.
2026-01-13 18:29:06 +08:00
godnight10061 0992657e54 Fix CI workflows for torch and default branch 2026-01-13 09:25:41 +08:00
godnight10061 2fa036992b Add unit regression for NaN oscillator mean 2026-01-13 08:56:46 +08:00
godnight10061 9432a34c38 Avoid MultivariateNormal crash on non-finite dynamics 2026-01-13 01:39:48 +08:00
scraed d4e8ee28fb Merge pull request #70 from godnight10061/fix/reshape-mask-3d
Thanks for the contribution! Tested and merged.
2026-01-11 18:22:14 +08:00
godnight10061 34b8e83831 Make CI tooling importable 2026-01-11 11:07:41 +08:00
godnight10061 e29d480a1a Fix reshape_mask for 3D masks 2026-01-10 14:46:14 +08:00
charrywhite 9479106c1d Fix Flux2 Dev Inpaint workflow size bug 2026-01-10 12:43:40 +08:00
charrywhite db4bbbe9f2 Update Flux.2.Dev workflow Preview pic 2026-01-10 11:47:29 +08:00
charrywhite 1a3e15964c Merge branch 'master' of https://github.com/scraed/LanPaint 2026-01-10 11:41:58 +08:00
charrywhite 4cc30f8a1d Fix Flux2.Dev Inpainting Workflow MaskBlend Bug 2026-01-10 11:41:53 +08:00
scraed d1e609192f Bump version from 1.4.8 to 1.4.9 2026-01-06 23:13:59 +08:00
charrywhite bf46831196 Update README.md 2026-01-06 23:01:04 +08:00
charrywhite c9d9a79b18 Update README.md 2026-01-06 22:59:21 +08:00
charrywhite 175156af86 Add Flux.2 Dev support 2026-01-06 22:57:42 +08:00
charrywhite 3281946b32 Merge branch 'master' of https://github.com/scraed/LanPaint 2026-01-06 22:42:16 +08:00
charrywhite 83f557ff4c Support Flux.2 Dev Inpainting 2026-01-06 22:42:07 +08:00
charrywhite 5e7fe4d5e4 Enhance README with new Discord info and features
Updated README to highlight new Discord channel and features.
2026-01-06 21:06:25 +08:00
scraed 69f2aded4d Update README with Discord link and new features
Added Discord link for community engagement and announced new inpainting and outpainting features.
2026-01-06 20:02:15 +08:00
scraed 5f1b6d8989 add discord link 2026-01-06 20:01:41 +08:00
scraed 7b0d144db9 Update Qwen Image Edit 2511 support
Updated references for Qwen Image Edit to include version 2511 alongside 2509.
2025-12-28 18:40:41 +08:00
scraed 27ebb6e7af fix qwen edit shape error 2025-12-25 23:18:51 +08:00
charrywhite f148e4b631 Update paper title in README.md 2025-12-17 22:33:46 +08:00
charrywhite cc9ec8873a Revise citation and add TMLR link
Updated citation format and added TMLR link.
2025-12-17 22:30:06 +08:00
scraed 565087e8f8 fix z image resize error for inpaint 2025-12-15 18:15:15 +08:00
scraed ebdc5dcc02 Merge branch 'master' of https://github.com/scraed/LanPaint 2025-12-06 10:26:06 +08:00
scraed 9336523576 more clear bug report 2025-12-06 10:26:01 +08:00
scraed 21520b38fb Update README.md 2025-12-05 09:39:13 +08:00
scraed 868bbaeefe fix detail display 2025-12-04 15:01:50 +08:00
scraed 3581d6c217 outpaint example 2025-12-04 14:58:16 +08:00
scraed d4da280a92 update version 2025-12-04 14:37:10 +08:00
scraed d0bcd459cc update version 2025-12-04 14:20:58 +08:00
scraed 469012d183 Z image update 2025-12-04 14:17:29 +08:00
scraed 3f4a5cdb5a update pyproject.yaml 2025-11-23 01:11:08 +08:00
scraed e0ce52b1be fix blur issue in Wan 2.2 14B 2025-11-23 00:46:25 +08:00
scraed 2c8f4baec0 add wan 2.2 5B 2025-11-17 16:09:45 +08:00
scraed daf1bdf4be update version 2025-11-17 16:05:42 +08:00
scraed f4d4dfdd61 Merge branch 'master' of https://github.com/scraed/LanPaint 2025-11-17 16:04:51 +08:00
scraed 6d624e7e1e support Wan 2.2 5B 2025-11-17 16:04:48 +08:00
scraed e0940059e6 Update README.md 2025-11-10 13:37:49 +08:00
scraed 76ce20b51f version change 2025-10-26 23:29:23 +08:00
scraed 29ed0b3da5 Add **kwargs to CFGGuider_LanPaint sampling methods
Updated the outer_sample and inner_sample method calls in CFGGuider_LanPaint to accept and forward arbitrary keyword arguments (**kwargs). This allows for greater flexibility and compatibility with future changes or additional parameters.
2025-10-26 23:28:35 +08:00
scraed 3d5aa3a3fa Add Hunyuan supports
Added new example workflows for Hunyuan, including images and JSON workflow files. Updated README with Hunyuan T2I inpainting instructions and links. Increased default LanPaint step size to 0.2 in nodes and advanced sampler, and refactored scale_latent_inpaint logic in lanpaint.py for improved compatibility.
2025-10-16 15:59:13 +08:00
charrywhite 6109df6591 Upload a 41frames outpaint workflows 2025-10-03 22:32:19 +08:00
scraed 536716a7b9 typo fix 2025-10-03 09:59:01 +08:00
scraed ece228ab10 Bump version to 1.4.0 in pyproject.toml
Updated the project version from 1.3.3 to 1.4.0 to prepare for a new release.
2025-10-02 23:24:21 +08:00
charrywhite c048230250 Update Original Video (Fix to 40 frames, fps=12) 2025-10-02 21:48:44 +08:00
charrywhite aa885d9fa6 Update README.md 2025-10-02 21:33:37 +08:00
charrywhite 80d55a47b0 Add 81 frames example 2025-10-02 21:31:19 +08:00
charrywhite 79db1765e3 Add video inpainting (Beta) 2025-10-02 21:13:43 +08:00
charrywhite 929328d08d Update README.md 2025-09-27 00:42:43 +08:00
scraed c12897040b Update README with inpainting result image
Replaced the workflow & masks link with an inpainting result image to better showcase the output of the project.
2025-09-24 18:49:34 +08:00
scraed 48b9a59ca3 Update README with new model support and examples
Added references to Wan2.2 and Qwen Image Edit 2509 support, and included a link to workflow and mask examples for improved documentation.
2025-09-24 18:48:46 +08:00
scraed d55f8c61d0 Add Qwen Image Edit 2509 example to README
Introduced a new section for Qwen Image Edit 2509 inpainting example, including workflow and model download links, and an example image. Updated the table of contents to reference the new example.
2025-09-24 18:43:41 +08:00
scraed a097b9c9b2 Update LanPaintQwen_04.jpg example image
Replaces the existing LanPaintQwen_04.jpg in the examples directory with a new version. This may reflect updated content or improved image quality.
2025-09-24 18:40:25 +08:00
scraed 70cc94cc6a update readme 2025-09-24 18:30:04 +08:00
scraed e653f45103 Support 2509 2025-09-24 18:27:11 +08:00
scraed c9c237c2de Add WAN 2.2 partial inpaint workflow template 2025-09-18 16:55:48 +08:00
scraed c5aa0116d5 Add Example 16 for partial inpainting with reference
Added a new example (Example 16) demonstrating partial inpainting using reference images. Updated README with instructions and links for the new example, and added related images to the examples directory.
2025-09-18 15:40:22 +08:00
75 changed files with 13801 additions and 168 deletions
+1
View File
@@ -26,6 +26,7 @@ jobs:
run: |
python -m pip install --upgrade pip
pip install .[dev]
pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
- name: Run Linting
run: |
ruff check .
+2
View File
@@ -11,3 +11,5 @@ jobs:
runs-on: ubuntu-latest
steps:
- uses: comfy-org/node-diff@main
with:
base_ref: ${{ github.event.repository.default_branch }}
+2
View File
@@ -98,3 +98,5 @@ cookiecutter-pypackage-env/
# vscode settings
.history/
*.code-workspace
.vscode/
/.vscode
+4
View File
@@ -5,4 +5,8 @@
"/PATH/TO/ComfyUI/",
"/PATH/TO/ComfyUI/custom_nodes/"
],
"cursorpyright.analysis.extraPaths": [
"/PATH/TO/ComfyUI/",
"/PATH/TO/ComfyUI/custom_nodes/"
],
}
+304 -20
View File
@@ -1,49 +1,115 @@
<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)
[![Blog](https://img.shields.io/badge/📝-Blog-9cf)](https://scraed.github.io/scraedBlog/)
[![GitHub stars](https://img.shields.io/github/stars/scraed/LanPaint)](https://github.com/scraed/LanPaint/stargazers)
[![Discord](https://img.shields.io/badge/Discord-5865F2?style=for-the-badge&logo=discord&logoColor=white)](https://discord.gg/yN5wYDE6W4)
</div>
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). 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.
![Qwen Result 2](https://github.com/scraed/LanPaint/blob/master/examples/LanPaintQwen_03.jpg)
Check [Mased Qwen Edit Workflow](https://github.com/scraed/LanPaint/tree/master/examples/Example_14). You need to follow the ComfyUI version of [Qwen Image Edit workflow](https://docs.comfy.org/tutorials/image/qwen/qwen-image-edit) to download and install the model.
The repository is for ComfyUI extension.
![Qwen Result 1](https://github.com/scraed/LanPaint/blob/master/examples/LanPaintQwen_01.jpg)
Diffusers Support: [LanPaint-Diffusers](https://github.com/charrywhite/LanPaint-diffusers) by [@charrywhite](https://github.com/charrywhite/)
Also check [Qwen Inpaint Workflow](https://github.com/scraed/LanPaint/tree/master/examples/Example_13) and [Qwen Outpaint Workflow](https://github.com/scraed/LanPaint/tree/master/examples/Example_12). You need to follow the ComfyUI version of [Qwen Image workflow](https://docs.comfy.org/tutorials/image/qwen/qwen-image) to download and install the model.
Benchmark code for paper reproduce: [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 2026: Join our discord!**
[Join our Discord](https://discord.gg/yN5wYDE6W4) to share experiences, discuss features, and explore future development.
`v1.5.0` fixes an important hidden bug that reduced performance and could blur images (especially with `z-image-base`) and also boosts overall LanPaint performance across other models.
If your inpainting results have wierd (glowing / broken) mask boundary, check this [issue](https://github.com/scraed/LanPaint/issues/80).
**🎬 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 Z-Image-Base too!**
| Original | Masked | Inpainted |
|:--------:|:------:|:---------:|
| ![Original Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/Original_No_Mask.png) | ![Masked Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/Masked_Load_Me_in_Loader.png) | ![Inpainted Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/InPainted_Drag_Me_to_ComfyUI.png) |
**🎬 NEW: LanPaint now supports video inpainting and outpainting based on Wan 2.2!**
<div align="center">
| Original Video | Mask (edit T-shirt text) | Inpainted Result |
|:--------------:|:----:|:----------------:|
| ![Original](https://github.com/scraed/LanPaint/blob/master/examples/Original_No_Mask-example18.gif) | ![Mask](https://github.com/scraed/LanPaint/blob/master/examples/Example_18/Masked_Load_Me_in_Loader.png) | ![Result](https://github.com/scraed/LanPaint/blob/master/examples/Inpainted_81frames_Drag_Me_to_ComfyUI_example18.gif) |
*Video Inpainting Example: 81 frames with temporal consistency*
</div>
Check our latest [Wan 2.2 Video Examples](#video-examples-beta), [Wan 2.2 Image Examples](#example-wan22-inpaintlanpaint-k-sampler-5-steps-of-thinking), and
[Qwen Image Edit 2509](#example-qwen-edit-2509-inpaint) support.
## Table of Contents
- [Features](#features)
- [Quickstart](#quickstart)
- [How to Use Examples](#how-to-use-examples)
- [Examples](#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)
- [Flux.2.Dev](#example-flux2dev-inpaintlanpaint-k-sampler-5-steps-of-thinking)
- [Flux 2 klein](#example-flux-2-klein-inpaintlanpaint-k-sampler-2-steps-of-thinking)
- [Z-image](#example-z-image-inpaintlanpaint-k-sampler-5-steps-of-thinking)
- [Z-image-base](#example-z-image-base-inpaintlanpaint-k-sampler-3-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)
- [Qwen Image Edit 2511 2509](#example-qwen-edit-2509-inpaint)
- [Qwen Image Edit 2508](#example-qwen-edit-2508-inpaint)
- [Qwen Image](#example-qwen-image-inpaintlanpaint-k-sampler-5-steps-of-thinking)
- [HiDream](#example-hidream-inpaint-lanpaint-k-sampler-5-steps-of-thinking)
- [SD 3.5](#example-sd-35-inpaintlanpaint-k-sampler-5-steps-of-thinking)
- [Flux](#example-flux-inpaintlanpaint-k-sampler-5-steps-of-thinking)
- [SDXL Examples](#example-sdxl-0-character-consistency-side-view-generation-lanpaint-k-sampler-5-steps-of-thinking)
- [SDXL](#example-sdxl-0-character-consistency-side-view-generation-lanpaint-k-sampler-5-steps-of-thinking)
- [Usage](#usage)
- [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)
## Features
- **Universal Compatibility** – Works instantly with almost any model (**SD 1.5, XL, 3.5, Flux, HiDream, Qwen-Image or custom LoRAs**) and ControlNet.
- **Universal Compatibility** – Works instantly with almost any model (**Z-image, Z-image-base, Hunyuan, Wan 2.2, Qwen Image/Edit, HiDream, SD 3.5, Flux-series, SDXL, SD 1.5 or custom LoRAs**) and ControlNet.
![Inpainting Result 13](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_13.jpg)
- **No Training Needed** – Works out of the box with your existing model.
- **Easy to Use** – Same workflow as standard ComfyUI KSampler.
@@ -79,7 +145,116 @@ Once installed, you'll find the LanPaint nodes under the "sampling" category in
- **[VAE Encode for Inpainting](https://comfyanonymous.github.io/ComfyUI_examples/inpaint/)**
- **[Set Latent Noise Mask](https://comfyui-wiki.com/en/tutorial/basic/how-to-inpaint-an-image-in-comfyui)**
## Examples
## Video Examples (Beta)
LanPaint now supports video inpainting with Wan 2.2, enabling you to seamlessly inpaint masked regions across video frames while maintaining temporal consistency.
**Note:** LanPaint supports video inpainting for longer sequences (e.g., 81 frames), but processing time increases significantly (please check the [Resource Consumption](#resource-consumption) section for details) and performance may become unstable. For optimal results and stability, we recommend limiting video inpainting to **40 frames or fewer**.
### Wan 2.2 Video Inpainting
*Example: Wan2.2 t2v 14B, 480p video (11:6), 40 frames, LanPaint K Sampler, 2 steps of thinking*
| Original Video | Mask (Add a white hat) | Inpainted Result |
|:--------------:|:----:|:----------------:|
| ![Original Video](https://github.com/scraed/LanPaint/blob/master/examples/Original_No_Mask_example17.gif) | ![Mask](https://github.com/scraed/LanPaint/blob/master/examples/Example_17/Masked_Load_Me_in_Loader.png) | ![Inpainted Result](https://github.com/scraed/LanPaint/blob/master/examples/Inpainted_40frames_Drag_Me_to_ComfyUI_example17.gif) |
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_17)
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.
*Example: Wan2.2 t2v 14B, 480p video (1:1 outpaint to 11:6), 40 frames, LanPaint K Sampler, 2 steps of thinking*
| Original Video | Mask (Expand to 880x480) | Outpainted Result |
|:--------------:|:----:|:-----------------:|
| ![Original Video](https://github.com/scraed/LanPaint/blob/master/examples/Original_Load_Me_in_Loader_example19.gif) | ![Mask](https://github.com/scraed/LanPaint/blob/master/examples/Mask_Example19_.png) | ![Outpainted Result](https://github.com/scraed/LanPaint/blob/master/examples/Outpainted_40frames_Drag_Me_to_ComfyUI_example19.gif) |
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_19)
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.
### Resource Consumption
<table>
<thead>
<tr>
<th align="left">Processing Mode</th>
<th align="left">Resolution</th>
<th align="left">Frames Processed</th>
<th align="left">VRAM Required</th>
<th align="left">Total Runtime (20 steps)</th>
</tr>
</thead>
<tbody>
<tr style="background-color: #e8f4f8;">
<td><strong>Inpainting</strong></td>
<td>880×480 (11:6)</td>
<td>40 frames</td>
<td>39.8 GB</td>
<td><strong>05:37 min</strong></td>
</tr>
<tr style="background-color: #e8f4f8;">
<td><strong>Inpainting</strong></td>
<td>480×480 (1:1)</td>
<td>40 frames</td>
<td>38.0 GB</td>
<td><strong>05:35 min</strong></td>
</tr>
<tr style="background-color: #e8f4f8;">
<td><strong>Outpainting</strong></td>
<td>880×480 (11:6)</td>
<td>40 frames</td>
<td>40.2 GB</td>
<td><strong>05:36 min</strong></td>
</tr>
<tr style="background-color: #fff4e6;">
<td><strong>Inpainting</strong></td>
<td>880×480 (11:6)</td>
<td>81 frames</td>
<td>43.3 GB</td>
<td><strong>16:23 min</strong></td>
</tr>
<tr style="background-color: #fff4e6;">
<td><strong>Inpainting</strong></td>
<td>480×480 (1:1)</td>
<td>81 frames</td>
<td>39.8 GB</td>
<td><strong>14:25 min</strong></td>
</tr>
<tr style="background-color: #fff4e6;">
<td><strong>Outpainting</strong></td>
<td>880×480 (11:6)</td>
<td>81 frames</td>
<td>42.6 GB</td>
<td><strong>13:46 min</strong></td>
</tr>
</tbody>
</table>
<sub>**Test Platform**: All tests were conducted on an NVIDIA RTX Pro 6000.<br>
**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:** 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
### Example Hunyuan T2I: InPaint(LanPaint K Sampler, 5 steps of thinking)
We are excited to announce that LanPaint now supports inpainting with Hunyuan text to image generation.
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_20)
You need to follow the ComfyUI version of [Hunyuan workflow](https://docs.comfy.org/tutorials/video/hunyuan-video#hunyuan-text-to-video-workflow) to download and install the model.
### Example Wan2.2: InPaint(LanPaint K Sampler, 5 steps of thinking)
We are excited to announce that LanPaint now supports Wan2.2 text to image generation with Wan2.2 T2V model.
@@ -90,6 +265,71 @@ 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 Z-image-base: InPaint(LanPaint K Sampler, 3 steps of thinking)
LanPaint also supports inpainting with the Z-image-base model.
**Warning (stability)**: Z-image-base can easily diverge with LanPaint. Start with **small `LanPaint_StepSize`** and **fewer thinking iterations** (lower `LanPaint_NumSteps`) and increase gradually only if stable.
<details open>
<summary>View Original / Masked / Inpainted Comparison</summary>
| Original | Masked | Inpainted |
|:--------:|:------:|:---------:|
| ![Original Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/Original_No_Mask.png) | ![Masked Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/Masked_Load_Me_in_Loader.png) | ![Inpainted Z-image-base](https://github.com/scraed/LanPaint/blob/master/examples/Example_25/InPainted_Drag_Me_to_ComfyUI.png) |
</details>
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_25)
Workflow template (JSON): [Z_image_base_Inpaint.json](https://github.com/scraed/LanPaint/blob/master/example_workflows/Z_image_base_Inpaint.json)
### 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.
![Inpainting Result 46](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_46.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_16)
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). This workflow also supports Qwen Image Edit 2511.
![Qwen Result 3](https://github.com/scraed/LanPaint/blob/master/examples/LanPaintQwen_04.jpg)
### Example Qwen Edit 2508: InPaint
![Qwen Result 2](https://github.com/scraed/LanPaint/blob/master/examples/LanPaintQwen_03.jpg)
Check [Mased Qwen Edit Workflow](https://github.com/scraed/LanPaint/tree/master/examples/Example_14). You need to follow the ComfyUI version of [Qwen Image Edit workflow](https://docs.comfy.org/tutorials/image/qwen/qwen-image-edit) to download and install the model.
### Example Qwen Image: InPaint(LanPaint K Sampler, 5 steps of thinking)
@@ -101,6 +341,9 @@ You need to follow the ComfyUI version of [Qwen Image workflow](https://docs.com
The following examples utilize a random seed of 0 to generate a batch of 4 images for variance demonstration and fair comparison. (Note: Generating 4 images may exceed your GPU memory; please adjust the batch size as necessary.)
![Qwen Result 1](https://github.com/scraed/LanPaint/blob/master/examples/LanPaintQwen_01.jpg)
Also check [Qwen Inpaint Workflow](https://github.com/scraed/LanPaint/tree/master/examples/Example_13) and [Qwen Outpaint Workflow](https://github.com/scraed/LanPaint/tree/master/examples/Example_12). You need to follow the ComfyUI version of [Qwen Image workflow](https://docs.comfy.org/tutorials/image/qwen/qwen-image) to download and install the model.
### Example HiDream: InPaint (LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 8](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_11.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_8)
@@ -119,6 +362,39 @@ You need to follow the ComfyUI version of [HiDream workflow](https://docs.comfy.
You need to follow the ComfyUI version of [SD 3.5 workflow](https://comfyui-wiki.com/en/tutorial/advanced/stable-diffusion-3-5-comfyui-workflow) to download and install the model.
### Example Flux.2.Dev: InPaint(LanPaint K Sampler, 5 steps of thinking)
<details open>
<summary>View Original / Masked / Inpainted Comparison</summary>
| Original | Masked | Inpainted |
|:--------:|:------:|:---------:|
| ![Original Flux.2.Dev](https://github.com/scraed/LanPaint/blob/master/examples/Example_23/Original_No_Mask.png) | ![Masked Flux.2.Dev](https://github.com/scraed/LanPaint/blob/master/examples/Example_23/Masked_Load_Me_in_Loader.png) | ![Inpainted Flux.2.Dev](https://github.com/scraed/LanPaint/blob/master/examples/Example_23/InPainted_Drag_Me_to_ComfyUI.png) |
</details>
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_23)
[Model Used in This Example](https://huggingface.co/Comfy-Org/flux2-dev)
(Note: Prompt First mode is disabled on Flux.2.Dev. As it does not use CFG guidance.)
### Example Flux 2 klein: InPaint(LanPaint K Sampler, 2 steps of thinking)
<details open>
<summary>View Original / Masked / Inpainted Comparison</summary>
| Original | Masked | Inpainted |
|:--------:|:------:|:---------:|
| ![Original Flux 2 klein](https://github.com/scraed/LanPaint/blob/master/examples/Example_24/Original_No_Mask.png) | ![Masked Flux 2 klein](https://github.com/scraed/LanPaint/blob/master/examples/Example_24/Masked_Load_Me_in_Loader.png) | ![Inpainted Flux 2 klein](https://github.com/scraed/LanPaint/blob/master/examples/Example_24/InPainted_Drag_Me_to_ComfyUI.png) |
</details>
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_24)
[Model Used in This Example](https://docs.comfy.org/zh-CN/tutorials/flux/flux-2-klein). If you have quality problem on Comfy 0.11 and 0.12, check [this issue](https://github.com/scraed/LanPaint/issues/80).
### Example Flux: InPaint(LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 7](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_10.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_7)
@@ -235,7 +511,14 @@ 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
- 2026/03/02
- `v1.5.0`: Fixed a hidden bug that hurt performance and caused image blur (especially on `z-image-base`), and improved overall LanPaint performance on other models too.
- 2026/01/30
- Add Z-image-base documentation and Example_25 workflow images.
- 2025/08/08
- Add Qwen image support
- 2025/06/21
@@ -263,17 +546,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={}
}
```
@@ -281,4 +566,3 @@ Submit a PR to add your tutorial/video here, or open an [Issue](https://github.c
+80 -2
View File
@@ -10,7 +10,85 @@ __author__ = """LanPaint"""
__email__ = "czhengac@connect.ust.hk"
__version__ = "0.0.1"
from .src.LanPaint.nodes import NODE_CLASS_MAPPINGS
from .src.LanPaint.nodes import NODE_DISPLAY_NAME_MAPPINGS
def _install_lightweight_runtime_stubs() -> None:
"""Install lightweight stubs so tooling can import this package without ComfyUI.
This is used by CI tooling (e.g., comfy-org/node-diff) that imports NODE_CLASS_MAPPINGS
in an environment where ComfyUI isn't installed.
"""
import sys
import types
# `src/LanPaint/nodes.py` uses `torch.Tensor` in type annotations.
try:
import torch # noqa: F401
except ModuleNotFoundError:
torch_mod = types.ModuleType("torch")
class Tensor: # noqa: N801 (match torch naming)
pass
torch_mod.Tensor = Tensor
torch_mod.nn = types.SimpleNamespace(functional=types.SimpleNamespace())
sys.modules["torch"] = torch_mod
if "comfyui_version" not in sys.modules:
comfyui_version_mod = types.ModuleType("comfyui_version")
comfyui_version_mod.__version__ = "0.0.0"
sys.modules["comfyui_version"] = comfyui_version_mod
sys.modules.setdefault("nodes", types.ModuleType("nodes"))
sys.modules.setdefault("latent_preview", types.ModuleType("latent_preview"))
if "comfy" not in sys.modules:
comfy_mod = types.ModuleType("comfy")
comfy_mod.__path__ = []
comfy_utils_mod = types.ModuleType("comfy.utils")
def repeat_to_batch_size(tensor, batch_size): # type: ignore[no-untyped-def]
if getattr(tensor, "shape", ())[0] == batch_size:
return tensor
return tensor
comfy_utils_mod.repeat_to_batch_size = repeat_to_batch_size
comfy_samplers_mod = types.ModuleType("comfy.samplers")
class DummyKSAMPLER: # noqa: N801 (match ComfyUI naming)
pass
comfy_samplers_mod.KSAMPLER = DummyKSAMPLER
comfy_model_base_mod = types.ModuleType("comfy.model_base")
class ModelType: # noqa: N801 (match ComfyUI naming)
FLUX = "FLUX"
FLOW = "FLOW"
class WAN22: # noqa: N801 (match ComfyUI naming)
pass
comfy_model_base_mod.ModelType = ModelType
comfy_model_base_mod.WAN22 = WAN22
comfy_mod.utils = comfy_utils_mod
comfy_mod.samplers = comfy_samplers_mod
comfy_mod.model_base = comfy_model_base_mod
sys.modules["comfy"] = comfy_mod
sys.modules["comfy.utils"] = comfy_utils_mod
sys.modules["comfy.samplers"] = comfy_samplers_mod
sys.modules["comfy.model_base"] = comfy_model_base_mod
try:
from .src.LanPaint.nodes import NODE_CLASS_MAPPINGS
from .src.LanPaint.nodes import NODE_DISPLAY_NAME_MAPPINGS
except ModuleNotFoundError:
_install_lightweight_runtime_stubs()
from .src.LanPaint.nodes import NODE_CLASS_MAPPINGS
from .src.LanPaint.nodes import NODE_DISPLAY_NAME_MAPPINGS
WEB_DIRECTORY = "./web"
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

File diff suppressed because it is too large Load Diff
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

File diff suppressed because it is too large Load Diff
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.3 MiB

+786
View File
@@ -0,0 +1,786 @@
{
"id": "978d3a45-3d13-43c6-8ef9-89dc3e74d6ba",
"revision": 0,
"last_node_id": 82,
"last_link_id": 220,
"nodes": [
{
"id": 66,
"type": "SetLatentNoiseMask",
"pos": [
514.1915893554688,
781.9396362304688
],
"size": [
264.5999755859375,
46
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 216
},
{
"name": "mask",
"type": "MASK",
"link": 189
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"slot_index": 0,
"links": [
186
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.23",
"Node name for S&R": "SetLatentNoiseMask"
},
"widgets_values": []
},
{
"id": 78,
"type": "CLIPTextEncode",
"pos": [
314.8565368652344,
255.63235473632812
],
"size": [
422.84503173828125,
164.31304931640625
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"label": "clip",
"name": "clip",
"type": "CLIP",
"link": 194
}
],
"outputs": [
{
"label": "CONDITIONING",
"name": "CONDITIONING",
"type": "CONDITIONING",
"slot_index": 0,
"links": [
195
]
}
],
"title": "CLIP Text Encode (Positive Prompt)",
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.26",
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
"cute anime girl with massive fluffy fennec ears and a big fluffy tail blonde messy long hair blue eyes wearing a maid outfit with a long black gold leaf pattern dress and a white apron mouth open placing a fancy black forest cake with candles on top of a dinner table of an old dark Victorian mansion lit by candlelight with a bright window to the foggy forest and very expensive stuff everywhere there are paintings on the walls"
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 80,
"type": "CLIPTextEncode",
"pos": [
362.7684020996094,
481.0662536621094
],
"size": [
422.84503173828125,
164.31304931640625
],
"flags": {
"collapsed": true
},
"order": 3,
"mode": 0,
"inputs": [
{
"label": "clip",
"name": "clip",
"type": "CLIP",
"link": 196
}
],
"outputs": [
{
"label": "CONDITIONING",
"name": "CONDITIONING",
"type": "CONDITIONING",
"slot_index": 0,
"links": [
210
]
}
],
"title": "CLIP Text Encode (Negative Prompt)",
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.26",
"Node name for S&R": "CLIPTextEncode"
},
"widgets_values": [
""
],
"color": "#322",
"bgcolor": "#533"
},
{
"id": 79,
"type": "FluxGuidance",
"pos": [
529.1380615234375,
156.20236206054688
],
"size": [
211.60000610351562,
58
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"label": "conditioning",
"name": "conditioning",
"type": "CONDITIONING",
"link": 195
}
],
"outputs": [
{
"label": "CONDITIONING",
"name": "CONDITIONING",
"type": "CONDITIONING",
"slot_index": 0,
"links": [
207
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.26",
"Node name for S&R": "FluxGuidance"
},
"widgets_values": [
3.5
]
},
{
"id": 77,
"type": "CheckpointLoaderSimple",
"pos": [
-66.25785064697266,
484.081787109375
],
"size": [
315,
98
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"label": "MODEL",
"name": "MODEL",
"type": "MODEL",
"slot_index": 0,
"links": [
197
]
},
{
"label": "CLIP",
"name": "CLIP",
"type": "CLIP",
"slot_index": 1,
"links": [
194,
196
]
},
{
"label": "VAE",
"name": "VAE",
"type": "VAE",
"slot_index": 2,
"links": [
200,
203
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.26",
"Node name for S&R": "CheckpointLoaderSimple"
},
"widgets_values": [
"flux1-dev-fp8.safetensors"
]
},
{
"id": 73,
"type": "LanPaint_KSampler",
"pos": [
1005.0399780273438,
296.9553527832031
],
"size": [
413.6495666503906,
596
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 197
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 207
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 210
},
{
"name": "latent_image",
"type": "LATENT",
"link": 186
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"slot_index": 0,
"links": [
187
]
}
],
"properties": {
"cnr_id": "LanPaint",
"ver": "56bd6c04e89124cd06682b304245d6ddf8b20522",
"Node name for S&R": "LanPaint_KSampler"
},
"widgets_values": [
0,
"fixed",
30,
1,
"euler",
"simple",
1,
5,
"Image First",
"LanPaint KSampler. Recommend steps 50, LanPaint NumSteps 1-20 depending on the difficulty of task. LanPaint_EndSigma = 3.0 for anime style, 0.6 for realistic style. For more information, visit https://github.com/scraed/LanPaint",
"🖼️ Image Inpainting"
]
},
{
"id": 65,
"type": "VAEEncode",
"pos": [
200.1034698486328,
776.3561401367188
],
"size": [
210,
46
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "pixels",
"type": "IMAGE",
"link": 188
},
{
"name": "vae",
"type": "VAE",
"link": 203
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"slot_index": 0,
"links": [
216
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.23",
"Node name for S&R": "VAEEncode"
},
"widgets_values": []
},
{
"id": 48,
"type": "SaveImage",
"pos": [
1075.49169921875,
1167.2703857421875
],
"size": [
311.2532653808594,
484.7096252441406
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 103
}
],
"outputs": [],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.23"
},
"widgets_values": [
"ComfyUI"
]
},
{
"id": 81,
"type": "LanPaint_MaskBlend",
"pos": [
1773.0189208984375,
1222.9478759765625
],
"size": [
210,
98
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "image1",
"type": "IMAGE",
"link": 219
},
{
"name": "image2",
"type": "IMAGE",
"link": 218
},
{
"name": "mask",
"type": "MASK",
"link": 220
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
217
]
}
],
"properties": {
"cnr_id": "LanPaint",
"ver": "4d3d5d17f0105b673df92da5b084cce567c9c712",
"Node name for S&R": "LanPaint_MaskBlend"
},
"widgets_values": [
9
]
},
{
"id": 82,
"type": "SaveImage",
"pos": [
2106.45166015625,
1090.682861328125
],
"size": [
311.2532653808594,
484.7096252441406
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 217
}
],
"outputs": [],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.23"
},
"widgets_values": [
"ComfyUI"
]
},
{
"id": 8,
"type": "VAEDecode",
"pos": [
1211.46484375,
1065.318359375
],
"size": [
210,
46
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 187
},
{
"name": "vae",
"type": "VAE",
"link": 200
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"slot_index": 0,
"links": [
103,
218
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.23",
"Node name for S&R": "VAEDecode"
},
"widgets_values": []
},
{
"id": 75,
"type": "LoadImage",
"pos": [
45.954593658447266,
1150.45556640625
],
"size": [
266.13720703125,
487.1314697265625
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
188,
219
]
},
{
"name": "MASK",
"type": "MASK",
"links": [
189,
220
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.27",
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"clipspace/clipspace-mask-8865503.5.png [input]",
"image"
]
}
],
"links": [
[
103,
8,
0,
48,
0,
"IMAGE"
],
[
186,
66,
0,
73,
3,
"LATENT"
],
[
187,
73,
0,
8,
0,
"LATENT"
],
[
188,
75,
0,
65,
0,
"IMAGE"
],
[
189,
75,
1,
66,
1,
"MASK"
],
[
194,
77,
1,
78,
0,
"CLIP"
],
[
195,
78,
0,
79,
0,
"CONDITIONING"
],
[
196,
77,
1,
80,
0,
"CLIP"
],
[
197,
77,
0,
73,
0,
"MODEL"
],
[
200,
77,
2,
8,
1,
"VAE"
],
[
203,
77,
2,
65,
1,
"VAE"
],
[
207,
79,
0,
73,
1,
"CONDITIONING"
],
[
210,
80,
0,
73,
2,
"CONDITIONING"
],
[
216,
65,
0,
66,
0,
"LATENT"
],
[
217,
81,
0,
82,
0,
"IMAGE"
],
[
218,
8,
0,
81,
1,
"IMAGE"
],
[
219,
75,
0,
81,
0,
"IMAGE"
],
[
220,
75,
1,
81,
2,
"MASK"
]
],
"groups": [
{
"id": 1,
"title": "Mask image for inpainting.",
"bounding": [
36.04227828979492,
989.7313232421875,
278.89093017578125,
669.3414916992188
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
},
{
"id": 2,
"title": "Convert Latents for LanPaint",
"bounding": [
286.0640563964844,
714.343505859375,
489.16796875,
197.81044006347656
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
},
{
"id": 3,
"title": "Load Model and Set Prompts",
"bounding": [
-78.9311294555664,
176.08712768554688,
1006.1721801757812,
514.258544921875
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
},
{
"id": 4,
"title": "Inpaint with the LanPaint KSampler",
"bounding": [
960.8922729492188,
179.17588806152344,
474.8909606933594,
630.4742431640625
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
},
{
"id": 5,
"title": "LanPaint OutPut",
"bounding": [
1085.6029052734375,
994.0775756835938,
345.4561767578125,
669.4969482421875
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
},
{
"id": 11,
"title": "LanPaint",
"bounding": [
-262.59381103515625,
140.46656799316406,
1737.328857421875,
797.4443359375
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"ds": {
"scale": 0.35049389948139237,
"offset": [
348.866804381099,
308.65057628971834
]
},
"frontendVersion": "1.27.10",
"node_versions": {
"comfy-core": "0.3.18",
"LanPaint": "0f509469ed2cd60c6032f739e282aad5dfc06166"
}
},
"version": 0.4
}
Binary file not shown.

After

Width:  |  Height:  |  Size: 670 KiB

File diff suppressed because it is too large Load Diff
Binary file not shown.

After

Width:  |  Height:  |  Size: 674 KiB

File diff suppressed because it is too large Load Diff
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

File diff suppressed because it is too large Load Diff
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

File diff suppressed because it is too large Load Diff
Binary file not shown.

After

Width:  |  Height:  |  Size: 687 KiB

File diff suppressed because it is too large Load Diff
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 221 KiB

Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 141 KiB

Binary file not shown.
Binary file not shown.
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.3 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.3 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.9 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.2 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.9 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.7 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.2 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 461 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 4.1 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 801 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 551 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 6.2 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 4.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 4.1 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 8.5 MiB

+4 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "LanPaint"
version = "1.3.2"
version = "1.5.2"
description = "Achieve seamless inpainting results without needing a specialized inpainting model."
authors = [
{name = "LanPaint", email = "czhengac@connect.ust.hk"}
@@ -75,5 +75,8 @@ select = [
# See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f
]
[tool.ruff.lint.per-file-ignores]
"src/LanPaint/nodes.py" = ["F403", "F405"]
[tool.ruff.lint.flake8-quotes]
inline-quotes = "double"
+337
View File
@@ -0,0 +1,337 @@
"""
Early Stop Logic Contributed by `https://github.com/godnight10061`.
"""
import inspect
from typing import Any, Callable, Optional
import torch
from .types import LangevinState
def _clamp01(val: float) -> float:
if val <= 0.0:
return 0.0
if val >= 1.0:
return 1.0
return val
def _abt_scale(abt_val: float) -> float:
"""
Smooth, parameter-free scale based on outer-step noise level.
- 0 at abt=0/1 (disable at extreme noise / extreme tail)
- 1 at abt=0.5 (mid-schedule)
"""
abt_val = _clamp01(abt_val)
return _clamp01(4.0 * abt_val * (1.0 - abt_val))
def _boundary_weight(latent_mask: torch.Tensor, inpaint_weight: torch.Tensor) -> Optional[torch.Tensor]:
"""
Return a 4-neighbor boundary weight: unknown pixels adjacent to known pixels.
This replaces the previous dilation-based "ring" (kernel/padding) and has no tunable hyperparameters.
"""
if latent_mask.dim() != 4:
return None
known = latent_mask > 0.5
neighbor_known = torch.zeros_like(known)
neighbor_known[:, :, 1:, :] |= known[:, :, :-1, :]
neighbor_known[:, :, :-1, :] |= known[:, :, 1:, :]
neighbor_known[:, :, :, 1:] |= known[:, :, :, :-1]
neighbor_known[:, :, :, :-1] |= known[:, :, :, 1:]
boundary = (~known) & neighbor_known
return boundary.to(dtype=torch.float32) * inpaint_weight
def _weighted_mse(t1: torch.Tensor, t2: torch.Tensor, weight: torch.Tensor) -> float:
diff_sq = (t1.to(dtype=torch.float32) - t2.to(dtype=torch.float32)) ** 2
denom = torch.sum(weight) + 1e-12
return float((torch.sum(diff_sq * weight) / denom).item())
class LanPaintEarlyStopper:
"""
Per-step early-stop logic for LanPaint inner (Langevin) iterations.
"""
@classmethod
def from_options(
cls,
*,
model_options: Optional[dict],
latent_mask: torch.Tensor,
abt: torch.Tensor,
default_threshold: float,
default_patience: int,
default_distance_fn: Optional[Callable[..., Any]],
) -> Optional["LanPaintEarlyStopper"]:
semantic_stop = model_options.get("lanpaint_semantic_stop") if isinstance(model_options, dict) else None
threshold = float(default_threshold)
patience = int(default_patience)
distance_fn = default_distance_fn
# distance_fn contract: return None (use default metric) or a scalar (Python number / 0-d (1-element) torch.Tensor)
if isinstance(semantic_stop, dict):
threshold = float(semantic_stop.get("threshold", threshold))
patience = int(semantic_stop.get("patience", patience))
distance_fn = semantic_stop.get("distance_fn", distance_fn)
# Backward compatibility: map legacy 'min_steps' to a patience floor so it is not an independent knob.
if patience > 0:
min_steps = semantic_stop.get("min_steps")
if min_steps is not None:
try:
min_steps_int = int(min_steps)
except (TypeError, ValueError):
min_steps_int = 0
if min_steps_int > 1:
patience = max(patience, min_steps_int - 1)
enabled_early_stop = (threshold > 0.0) and (patience > 0)
# Require N+1 consecutive stable checks:
# - the first stable step sets patience_counter to 1
# - `patience=1` therefore stops after 2 stable steps
patience_eff = max(1, patience) + 1
threshold_eff = threshold
inpaint_weight = ring_weight = trace = abt_val = None
if enabled_early_stop:
try:
abt_val = float(torch.mean(abt).item())
except (TypeError, ValueError):
abt_val = 0.0
threshold_eff = threshold * _abt_scale(abt_val)
if threshold_eff <= 0.0:
enabled_early_stop = False
else:
inpaint_weight = (1 - latent_mask).to(dtype=torch.float32)
if float(torch.sum(inpaint_weight).item()) < 1e-6:
enabled_early_stop = False
else:
ring_weight = _boundary_weight(latent_mask, inpaint_weight)
if isinstance(model_options, dict):
trace = model_options.get("lanpaint_semantic_trace")
if not enabled_early_stop:
return None
# Pre-fetch trace keys to avoid repeated dict lookups
bench_case_id = bench_outer_step = bench_timestep = None
if isinstance(trace, list) and isinstance(model_options, dict):
bench_case_id = model_options.get("bench_case_id")
bench_outer_step = model_options.get("bench_outer_step")
bench_timestep = model_options.get("bench_timestep")
return cls(
enabled=enabled_early_stop,
threshold=threshold,
threshold_eff=threshold_eff,
patience_eff=patience_eff,
inpaint_weight=inpaint_weight,
ring_weight=ring_weight,
distance_fn=distance_fn,
trace=trace,
bench_case_id=bench_case_id,
bench_outer_step=bench_outer_step,
bench_timestep=bench_timestep,
abt_val=abt_val,
)
def __init__(
self,
*,
enabled: bool,
threshold: float,
threshold_eff: float,
patience_eff: int,
inpaint_weight: Optional[torch.Tensor],
ring_weight: Optional[torch.Tensor],
distance_fn: Optional[Callable[..., Any]] = None,
trace: Optional[list] = None,
bench_case_id: Any = None,
bench_outer_step: Any = None,
bench_timestep: Any = None,
abt_val: Optional[float] = None,
) -> None:
self.enabled = bool(enabled)
self.threshold = float(threshold)
self.threshold_eff = float(threshold_eff)
self.patience_eff = int(patience_eff)
self.inpaint_weight = inpaint_weight
self.ring_weight = ring_weight
self.trace = trace
self.bench_case_id = bench_case_id
self.bench_outer_step = bench_outer_step
self.bench_timestep = bench_timestep
self.abt_val = abt_val
self.patience_counter = 0
self.x0_anchor = None
self._dist_wrapper = self._wrap_distance_fn(distance_fn) if self.enabled else None
@property
def has_custom_distance_fn(self) -> bool:
return self._dist_wrapper is not None
@staticmethod
def _wrap_distance_fn(distance_fn: Optional[Callable[..., Any]]):
"""
Wrap a user-provided `distance_fn` into a normalized callable: fn(prev, cur, ctx) -> dist|None.
Supported signatures:
- 3+ positional (or *args): `distance_fn(prev, cur, ctx)`
- explicit / **kwargs ctx: `distance_fn(prev, cur, ctx=ctx)`
- default 2-arg: `distance_fn(cur, prev)`
Return contract: None (use default metric) or a scalar (Python number / 0-d (1-element) torch.Tensor).
"""
if not callable(distance_fn):
return None
try:
sig = inspect.signature(distance_fn)
params = list(sig.parameters.values())
has_ctx_param = "ctx" in sig.parameters
has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params)
has_var_pos = any(p.kind == inspect.Parameter.VAR_POSITIONAL for p in params)
pos_params = [
p
for p in params
if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
]
if len(pos_params) >= 3 or has_var_pos:
# 3-arg positional: fn(prev, cur, ctx)
return lambda p, c, ctx: distance_fn(p, c, ctx)
if has_ctx_param or has_var_kw:
# keyword ctx: fn(prev, cur, ctx=ctx)
return lambda p, c, ctx: distance_fn(p, c, ctx=ctx)
# Default 2-arg: fn(cur, prev)
return lambda p, c, ctx: distance_fn(c, p)
except (ValueError, TypeError):
# Fallback for built-ins or complex callables.
def fallback_wrapper(p, c, ctx):
try:
return distance_fn(p, c, ctx)
except TypeError as e:
tb = e.__traceback__
if tb is not None and tb.tb_frame.f_code is not fallback_wrapper.__code__:
raise
return distance_fn(c, p)
return fallback_wrapper
def step(
self,
*,
i: int,
n_steps: int,
x_t_before: torch.Tensor,
x_t_after: torch.Tensor,
x_t_prev_for_custom: Optional[torch.Tensor],
prev_args: Any,
args: Any,
ctx: dict,
) -> bool:
if not self.enabled:
return False
# 'inpaint_weight' is guaranteed to be set when enabled is True in the caller.
inpaint = self.inpaint_weight
if inpaint is None:
return False
dist = None
custom_dist = False
dist_inpaint = dist_ring = dist_drift = x0_prev = x0_cur = None
if self._dist_wrapper is not None:
dist = self._dist_wrapper(x_t_prev_for_custom, x_t_after, ctx)
if dist is not None:
if isinstance(dist, torch.Tensor):
if dist.numel() != 1:
raise TypeError("distance_fn must return None or a scalar / 0-d (1-element) tensor")
dist = float(dist.item())
else:
dist = float(dist)
custom_dist = dist is not None
if dist is None:
def _get_x0(arg: Any) -> Optional[torch.Tensor]:
if isinstance(arg, LangevinState):
return arg.x0
if isinstance(arg, tuple) and len(arg) >= 3:
return arg[2]
return None
x0_prev = _get_x0(prev_args)
x0_cur = _get_x0(args)
if x0_prev is not None and x0_cur is not None:
dist_inpaint = _weighted_mse(x0_cur, x0_prev, inpaint)
dist_ring = _weighted_mse(x0_cur, x0_prev, self.ring_weight) if self.ring_weight is not None else None
dist = dist_inpaint if dist_ring is None else max(dist_inpaint, dist_ring)
else:
dist_inpaint = _weighted_mse(x_t_after, x_t_before, inpaint)
dist = dist_inpaint
threshold_used = self.threshold if custom_dist else self.threshold_eff
# Drift guard (only for default metric with x0_cur).
if x0_cur is not None and not custom_dist:
if dist <= threshold_used:
if self.x0_anchor is None:
self.x0_anchor = x0_cur.detach()
else:
drift_inpaint = _weighted_mse(x0_cur, self.x0_anchor, inpaint)
drift_ring = _weighted_mse(x0_cur, self.x0_anchor, self.ring_weight) if self.ring_weight is not None else None
dist_drift = drift_inpaint if drift_ring is None else max(drift_inpaint, drift_ring)
dist = max(dist, dist_drift)
else:
self.x0_anchor = None
if dist <= threshold_used:
self.patience_counter += 1
else:
self.patience_counter = 0
self.x0_anchor = None
should_stop = self.patience_counter >= self.patience_eff
if isinstance(self.trace, list):
self.trace.append(
{
"case_id": self.bench_case_id,
"outer_step": self.bench_outer_step,
"bench_timestep": self.bench_timestep,
"inner_step": i + 1,
"dist": dist,
"dist_inpaint": None if dist_inpaint is None else float(dist_inpaint),
"dist_ring": None if dist_ring is None else float(dist_ring),
"dist_drift": None if dist_drift is None else float(dist_drift),
"threshold": float(threshold_used),
"threshold_eff": float(self.threshold_eff),
"patience_counter": int(self.patience_counter),
"patience_eff": int(self.patience_eff),
"abt": None if self.abt_val is None else float(self.abt_val),
"custom_dist": bool(custom_dist),
"stopped": bool(should_stop),
}
)
return bool(should_stop)
+149 -34
View File
@@ -1,9 +1,11 @@
import torch
from .utils import *
from .utils import StochasticHarmonicOscillator
from functools import partial
from .earlystop import LanPaintEarlyStopper
from .types import LangevinState
class LanPaint():
def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False):
def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False, EarlyStopThreshold = 0.0, EarlyStopPatience = 1, EarlyStopHook = None):
self.n_steps = NSteps
self.chara_lamb = Lambda
self.IS_FLUX = IS_FLUX
@@ -13,6 +15,10 @@ class LanPaint():
self.friction = Friction
self.chara_beta = Beta
self.img_dim_size = None
self.early_stop_threshold = EarlyStopThreshold
self.early_stop_patience = EarlyStopPatience
self.early_stop_hook = EarlyStopHook
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)
@@ -21,23 +27,38 @@ class LanPaint():
# Create a tuple with ':' for the first dimension and 'None' repeated num_nones times
index = (slice(None),) + (0,) * (self.img_dim_size-1)
return array[index]
def unpack_model_output(self, output):
# Some guider/model wrappers return one denoised latent, others return
# both the normal and BIG-guidance denoised latents.
if isinstance(output, (tuple, list)):
if len(output) >= 2:
return output[0], output[1]
if len(output) == 1:
return output[0], output[0]
raise ValueError("Model output is empty")
return output, output
def __call__(self, x, latent_image, noise, sigma, latent_mask, current_times, model_options, seed, n_steps=None):
self.img_dim_size = len(x.shape)
self.latent_image = latent_image
self.noise = noise
if torch.mean(torch.abs(self.noise)) < 1e-8:
self.noise = torch.randn_like(self.noise)
if n_steps is None:
n_steps = self.n_steps
return self.LanPaint(x, sigma, latent_mask, current_times, n_steps, model_options, seed, self.IS_FLUX, self.IS_FLOW)
def LanPaint(self, x, sigma, latent_mask, current_times, n_steps, model_options, seed, IS_FLUX, IS_FLOW):
input_x = x
VE_Sigma, abt, Flow_t = current_times
step_size = self.step_size * (1 - abt)
step_size = self.add_none_dims(step_size)
# self.inner_model.inner_model.scale_latent_inpaint returns variance exploding x_t values
# This is the replace step
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
def scale_latent_inpaint(x, sigma, noise, latent_image):
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 )
else:
@@ -46,29 +67,74 @@ class LanPaint():
############ 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
stopper = LanPaintEarlyStopper.from_options(
model_options=model_options if isinstance(model_options, dict) else None,
latent_mask=latent_mask,
abt=abt,
default_threshold=self.early_stop_threshold,
default_patience=self.early_stop_patience,
default_distance_fn=self.early_stop_hook,
)
for i in range(n_steps):
score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = self.add_none_dims(abt), sigma = self.add_none_dims(VE_Sigma), tflow = self.add_none_dims(Flow_t), model_options = model_options, seed = seed )
prev_args = args
x_t_prev = x_t.detach() if (stopper is not None and stopper.has_custom_distance_fn) else None
x_t_before = x_t if (stopper is not None and stopper.enabled) else None
x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size , current_times, sigma_x = self.add_none_dims(self.sigma_x(abt)), sigma_y = self.add_none_dims(self.sigma_y(abt)), args = args)
if stopper is not None:
ctx = {
"step": i,
"steps_done": i + 1,
"n_steps": n_steps,
"mask": latent_mask,
"latent_image": self.latent_image,
"current_times": current_times,
"seed": seed,
}
if stopper.step(
i=i,
n_steps=n_steps,
x_t_before=x_t_before,
x_t_after=x_t,
x_t_prev_for_custom=x_t_prev,
prev_args=prev_args,
args=args,
ctx=ctx,
):
break
if IS_FLUX or IS_FLOW:
x = x_t / ( self.add_none_dims(abt)**0.5 + (1-self.add_none_dims(abt))**0.5 )
else:
x = x_t * ( 1+self.add_none_dims(VE_Sigma)**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, _ = self.unpack_model_output(
self.inner_model(x, sigma, model_options=model_options, seed=seed)
)
out = out * (1-latent_mask) + self.latent_image * latent_mask
input_x.copy_(x)
return out
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)
x_0, x_0_BIG = self.unpack_model_output(
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)
x_0, x_0_BIG = self.unpack_model_output(
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)
@@ -81,6 +147,13 @@ class LanPaint():
return beta
def langevin_dynamics(self, x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):
if args is not None and not isinstance(args, LangevinState):
if isinstance(args, tuple):
if len(args) == 2:
# Backwards compat: older state was (v, C) without x0.
args = LangevinState(args[0], args[1], None)
elif len(args) >= 3:
args = LangevinState(args[0], args[1], args[2])
# prepare the step size and time parameters
with torch.autocast(device_type=x_t.device.type, dtype=torch.float32):
step_sizes = self.prepare_step_size(current_times, step_size, sigma_x, sigma_y)
@@ -98,11 +171,10 @@ class LanPaint():
dt = dtx * (1-mask) + dty * mask
Gamma = Gamma_x * (1-mask) + Gamma_y * mask
def Coef_C(x_t):
x0 = self.x0_evalutation(x_t, score, sigma, args)
x0 = x_t + score(x_t)
C = (abt**0.5 * x0 - x_t )/ (1-abt) + A * x_t
return C
return C, x0
def advance_time(x_t, v, dt, Gamma, A, C, D):
dtype = x_t.dtype
with torch.autocast(device_type=x_t.device.type, dtype=torch.float32):
@@ -111,25 +183,74 @@ class LanPaint():
x_t = x_t.to(dtype)
v = v.to(dtype)
return x_t, v
if args is None:
#v = torch.zeros_like(x_t)
v = None
C = Coef_C(x_t)
#print(torch.squeeze(dtx), torch.squeeze(dty))
x_t, v = advance_time(x_t, v, dt, Gamma, A, C, D)
else:
v, C = args
x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
def advance_time_overdamped(x_t, dt, A, C, D):
"""
Overdamped (Gamma -> infinity) limit:
dx = -A x dt + C dt + D dW_t
with C treated as constant over this substep.
"""
dtype = x_t.dtype
with torch.autocast(device_type=x_t.device.type, dtype=torch.float32):
A_dt = A * dt
exp_neg = torch.exp(-A_dt)
C_new = Coef_C(x_t)
v = v + Gamma**0.5 * ( C_new - C) *dt
eps = 1e-8
abs_A = torch.abs(A)
# k = (1 - exp(-A dt)) / A -> dt when A -> 0
k = torch.where(abs_A < eps, dt, (-torch.expm1(-A_dt)) / A)
# k2 = (1 - exp(-2 A dt)) / (2 A) -> dt when A -> 0
k2 = torch.where(abs_A < eps, dt, (-torch.expm1(-2 * A_dt)) / (2 * A))
x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
mean = exp_neg * x_t + k * C
var = (D ** 2) * k2
noise = torch.randn_like(x_t) * torch.sqrt(torch.clamp(var, min=0.0))
x_t = mean + noise
return x_t.to(dtype)
C = C_new
return x_t, (v, C)
def run_damped(x_t, args):
if args is None:
v = None
C, x0 = Coef_C(x_t)
x_t, v = advance_time(x_t, v, dt, Gamma, A, C, D)
else:
v = args.v
C = args.C
x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
C_new, x0 = Coef_C(x_t)
v = v + Gamma**0.5 * ( C_new - C) *dt
x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
C = C_new
# args is (v, C, x0) for the next inner step.
return x_t, LangevinState(v, C, x0)
def run_overdamped(x_t, args):
if args is None:
C, x0 = Coef_C(x_t)
x_t = advance_time_overdamped(x_t, dt, A, C, D)
else:
C = args.C
x_t = advance_time_overdamped(x_t, dt / 2, A, C, D)
C_new, x0 = Coef_C(x_t)
x_t = x_t + (C_new - C) * dt
x_t = advance_time_overdamped(x_t, dt / 2, A, C, D)
C = C_new
# args is (v, C, x0); v is None in the overdamped fallback.
return x_t, LangevinState(None, C, x0)
try:
x_t_next, state = run_damped(x_t, args)
v_next = state.v
if torch.isnan(x_t_next).any() or (v_next is not None and torch.isnan(v_next).any()):
raise ValueError("NaN detected")
x_t = x_t_next
except Exception:
x_t, state = run_overdamped(x_t, args)
# args is (v, C, x0); v can be None if we fell back to the overdamped update.
return x_t, state
def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y):
# -------------------------------------------------------------------------
@@ -140,7 +261,7 @@ class LanPaint():
# Compute time step (dtx, dty) for x and y branches.
dtx = 2 * step_size * sigma_x
dty = 2 * step_size * sigma_y
# -------------------------------------------------------------------------
# Define friction parameter Gamma_hat for each branch.
# Using dtx**0 provides a tensor of the proper device/dtype.
@@ -165,9 +286,3 @@ class LanPaint():
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):
x0 = x_t + score(x_t)
return x0
+153 -70
View File
@@ -1,29 +1,87 @@
from contextlib import contextmanager
from inspect import cleandoc
import inspect
import math
# import nodes.py
import comfy
import nodes
import latent_preview
from functools import partial
import torch
from comfy.utils import repeat_to_batch_size
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 _version_tuple(value):
return tuple(int(part) if part.isdigit() else 0 for part in value.split("."))
def reshape_mask(input_mask, output_shape):
COMFYUI_VERSION_060_OR_NEWER = _version_tuple(comfyui_version.__version__) >= (0, 6, 0)
def reshape_mask(input_mask, output_shape,video_inpainting=False):
dims = len(output_shape) - 2
print('output shape',output_shape)
scale_mode = "nearest-exact"
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('input mask',input_mask.shape,type(input_mask),torch.max(input_mask),torch.min(input_mask))
print('target output_shape',output_shape)
print('input_mask.ndim:', input_mask.ndim, 'output_shape len:', len(output_shape))
# Handle input mask dimensions
if input_mask.ndim == 2:
input_mask = input_mask.unsqueeze(0).unsqueeze(0)
elif input_mask.ndim == 3:
input_mask = input_mask.unsqueeze(1)
# Handle 5D output shape (B, C, F, H, W) by ensuring input is 5D
if len(output_shape) == 5 and input_mask.ndim == 4:
if COMFYUI_VERSION_060_OR_NEWER:
input_mask = input_mask.unsqueeze(2) # (B, C, 1, H, W)
# Handle video case with temporal dimension
if video_inpainting: # Video case: (batch, channels, frames, height, width)
target_frames = output_shape[2]
target_height, target_width = output_shape[-2:]
print('Video case - input_mask initial shape:', input_mask.shape)
# First reshape input_mask to have proper dimensions for video processing
# Assume input is (frames, channels, height, width) -> (1, channels, frames, height, width)
## if comfy version < 0.6.0
if not COMFYUI_VERSION_060_OR_NEWER:
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
print('Video case - dimensions: batch_size={}, channels={}, frames={}, height={}, width={}'.format(batch_size, channels, frames, height, width))
print('Video case - target size:', (target_frames, target_height, target_width))
# 3D nearest-exact interpolation: (batch, channels, frames, height, width) -> (batch, channels, target_frames, target_height, target_width)
temp_mask = torch.nn.functional.interpolate(
input_mask,
size=(target_frames, target_height, target_width),
mode=scale_mode,
)
# temp_mask is already 5D: (batch, channels, target_frames, target_height, target_width)
mask = temp_mask
print('after mask',mask.shape)
# Handle channel dimension expansion if needed
if mask.shape[1] < output_shape[1]:
mask = mask.repeat(1, output_shape[1], 1, 1, 1)[:, :output_shape[1]]
# Handle batch dimension
mask = repeat_to_batch_size(mask, output_shape[0])
else: # Original 2D image case
if not COMFYUI_VERSION_060_OR_NEWER:
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])
return mask
def prepare_mask(noise_mask, shape, device):
return reshape_mask(noise_mask, shape).to(device)
def prepare_mask(noise_mask, shape, device,video_inpainting=False):
return reshape_mask(noise_mask, shape,video_inpainting).to(device)
def sampling_function_LanPaint(model, x, timestep, uncond, cond, cond_scale, cond_scale_BIG, model_options={}, seed=None):
if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
uncond_ = None
@@ -42,13 +100,18 @@ def sampling_function_LanPaint(model, x, timestep, uncond, cond, cond_scale, con
class CFGGuider_LanPaint:
def outer_sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
def outer_sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None, **kwargs):
print("CFGGuider outer_sample")
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:
denoise_mask = prepare_mask(denoise_mask, noise.shape, device)
video_inpainting = self.model_options.get("video_inpainting", False)
denoise_mask = prepare_mask(denoise_mask, noise.shape, device, video_inpainting)
noise = noise.to(device)
latent_image = latent_image.to(device)
@@ -57,7 +120,7 @@ class CFGGuider_LanPaint:
try:
self.model_patcher.pre_run()
output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, **kwargs)
finally:
self.model_patcher.cleanup()
@@ -86,14 +149,16 @@ 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 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
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
@@ -118,7 +183,7 @@ class KSamplerX0Inpaint:
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)
# Add TAESD preview support - directly use the latent_preview module
current_step = model_options.get("i", kwargs.get("i", 0))
total_steps = model_options.get("total_steps", 0)
@@ -129,7 +194,7 @@ class KSamplerX0Inpaint:
callback = model_options.get("callback", None)
if callback is not None:
callback({"i": current_step, "denoised": out, "x": x})
return out
# Custom sampler class extending ComfyUI's KSAMPLER for LanPaint
@@ -162,7 +227,10 @@ class KSAMPLER(comfy.samplers.KSAMPLER):
model_wrap.model_patcher.LanPaint_Beta,
model_wrap.model_patcher.LanPaint_StepSize,
IS_FLUX = IS_FLUX,
IS_FLOW = IS_FLOW)
IS_FLOW = IS_FLOW,
EarlyStopThreshold = getattr(model_wrap.model_patcher, "LanPaint_InnerThreshold", 0.0),
EarlyStopPatience = getattr(model_wrap.model_patcher, "LanPaint_InnerPatience", 1),
EarlyStopHook = extra_args.get("model_options", {}).get("lanpaint_semantic_hook", None))
model_k.LanPaint_early_stop = model_wrap.model_patcher.LanPaint_EarlyStop
#if not inpainting, after noise_scaling, noise = noise * sigma, which is the noise added to the clean latent image in the variance exploding diffusion model notation.
#if inpainting, after noise_scaling, noise = latent_image + noise * sigma, which is x_t in the variance exploding diffusion model notation for the known region.
@@ -245,6 +313,7 @@ class LanPaint_KSampler():
"LanPaint_NumSteps": ("INT", {"default": 5, "min": 0, "max": 100, "tooltip": "The number of steps for the Langevin dynamics, representing the turns of thinking per step."}),
"LanPaint_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: emphasis image quality, Prompt First: emphasis prompt following"}),
"LanPaint_Info": ("STRING", {"default": "LanPaint KSampler. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
"Inpainting_mode": (["🖼️ Image Inpainting", "🎬 Video Inpainting"], {"default": "🖼️ Image Inpainting", "tooltip": "Choose Image mode for photos or Video mode for video frames with temporal consistency"}),
}
}
@@ -255,18 +324,27 @@ class LanPaint_KSampler():
CATEGORY = "sampling"
DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image."
def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, LanPaint_NumSteps=5, LanPaint_PromptMode = "Image First", LanPaint_Info=""):
def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, LanPaint_NumSteps=5, LanPaint_PromptMode="Image First", LanPaint_Info="",Inpainting_mode="🖼️ Image Inpainting"):
model.LanPaint_StepSize = 0.15
model.LanPaint_StepSize = 0.2
model.LanPaint_Lambda = 16.0
model.LanPaint_Beta = 1.
model.LanPaint_NumSteps = LanPaint_NumSteps
model.LanPaint_Friction = 15.
model.LanPaint_EarlyStop = 1
model.LanPaint_InnerThreshold = 0.0
model.LanPaint_InnerPatience = 1
if LanPaint_PromptMode == "Image First":
model.LanPaint_cfg_BIG = cfg
else:
model.LanPaint_cfg_BIG = 0*cfg - 0.5
# Convert inpainting_mode to boolean for video_inpainting
video_inpainting = (Inpainting_mode == "🎬 Video Inpainting")
if not hasattr(model, 'model_options') or model.model_options is None:
model.model_options = {}
model.model_options["video_inpainting"] = video_inpainting
with override_sample_function():
return nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
class LanPaint_KSamplerAdvanced:
@@ -288,12 +366,15 @@ class LanPaint_KSamplerAdvanced:
"return_with_leftover_noise": (["disable", "enable"], ),
"LanPaint_NumSteps": ("INT", {"default": 5, "min": 0, "max": 100, "tooltip": "The number of steps for the Langevin dynamics, representing the turns of thinking per step."}),
"LanPaint_Lambda": ("FLOAT", {"default": 16., "min": 0.1, "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The bidirectional guidance scale. Higher values align with known regions more closely, but may result in instability."}),
"LanPaint_StepSize": ("FLOAT", {"default": 0.15, "min": 0.0001, "max": 1., "step": 0.01, "round": 0.001, "tooltip": "The step size for the Langevin dynamics. Higher values result in faster convergence but may be unstable."}),
"LanPaint_StepSize": ("FLOAT", {"default": 0.2, "min": 0.0001, "max": 1., "step": 0.01, "round": 0.001, "tooltip": "The step size for the Langevin dynamics. Higher values result in faster convergence but may be unstable."}),
"LanPaint_Beta": ("FLOAT", {"default": 1., "min": 0.0001, "max": 5, "step": 0.1, "round": 0.1, "tooltip": "The step size ratio between masked / unmasked regions. Lower value can compensate high values of LanPaint_Lambda."}),
"LanPaint_Friction": ("FLOAT", {"default": 15, "min": 0., "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The friction parameter for fast langevin, lower values result in faster convergence but may be unstable."}),
"LanPaint_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: emphasis image quality, Prompt First: emphasis prompt following"}),
"LanPaint_EarlyStop": ("INT", {"default": 1, "min": 0, "max": 10000, "tooltip": "The number of steps to stop the LanPaint early, useful for preventing the image from irregular patterns."}),
"LanPaint_Info": ("STRING", {"default": "LanPaint KSampler Adv. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
"Inpainting_mode": (["🖼️ Image Inpainting", "🎬 Video Inpainting"], {"default": "🖼️ Image Inpainting", "tooltip": "Choose Image mode for photos or Video mode for video frames with temporal consistency"}),
"LanPaint_InnerThreshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.0001, "round": 0.0001, "tooltip": "Early stop threshold for Langevin iterations based on semantic distance. 0.0 to disable. (Contributed by godnight10061)"}),
"LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop. (Contributed by godnight10061)"}),
},
}
@@ -302,7 +383,7 @@ class LanPaint_KSamplerAdvanced:
CATEGORY = "sampling"
def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0, LanPaint_StepSize=0.05, LanPaint_Lambda=5, LanPaint_Beta=1, LanPaint_NumSteps=5, LanPaint_Friction=5, LanPaint_PromptMode = "Image First", LanPaint_EarlyStop = 1, LanPaint_Info=""):
def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, LanPaint_NumSteps=5, LanPaint_Lambda=16.0, LanPaint_StepSize=0.2, LanPaint_Beta=1.0, LanPaint_Friction=15.0, LanPaint_PromptMode="Image First", LanPaint_EarlyStop=1, LanPaint_Info="", Inpainting_mode="🖼️ Image Inpainting", LanPaint_InnerThreshold=0.0, LanPaint_InnerPatience=1):
force_full_denoise = True
if return_with_leftover_noise == "enable":
force_full_denoise = False
@@ -315,13 +396,21 @@ class LanPaint_KSamplerAdvanced:
model.LanPaint_NumSteps = LanPaint_NumSteps
model.LanPaint_Friction = LanPaint_Friction
model.LanPaint_EarlyStop = LanPaint_EarlyStop
model.LanPaint_InnerThreshold = LanPaint_InnerThreshold
model.LanPaint_InnerPatience = LanPaint_InnerPatience
if LanPaint_PromptMode == "Image First":
model.LanPaint_cfg_BIG = cfg
else:
model.LanPaint_cfg_BIG = 0*cfg - 0.5
# Convert inpainting_mode to boolean for video_inpainting
video_inpainting = (Inpainting_mode == "🎬 Video Inpainting")
if not hasattr(model, 'model_options') or model.model_options is None:
model.model_options = {}
model.model_options["video_inpainting"] = video_inpainting
with override_sample_function():
return nodes.common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
return nodes.common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
class MaskBlend:
@@ -349,15 +438,23 @@ 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
kernel = self.gaussian_kernel(blend_overlap)
kernel = kernel.to(image1.device)
kernel = kernel[None, None, ...]
mask = torch.nn.functional.conv2d(mask[:,None,:,:], kernel, padding=blend_overlap//2)[:,0,:,:]
@@ -416,12 +513,14 @@ class LanPaint_SamplerCustom:
CATEGORY = "sampling/custom_sampling"
def sample(self, model, sampler, sigmas, add_noise, noise_seed, cfg, positive, negative, latent_image, LanPaint_NumSteps, LanPaint_PromptMode, LanPaint_Info=""):
model.LanPaint_StepSize = 0.15
model.LanPaint_StepSize = 0.2
model.LanPaint_Lambda = 16.0
model.LanPaint_Beta = 1.
model.LanPaint_NumSteps = LanPaint_NumSteps
model.LanPaint_Friction = 15.
model.LanPaint_EarlyStop = 1
model.LanPaint_InnerThreshold = 0.0
model.LanPaint_InnerPatience = 1
if LanPaint_PromptMode == "Image First":
model.LanPaint_cfg_BIG = cfg
else:
@@ -455,41 +554,37 @@ class LanPaint_SamplerCustom:
else:
out_denoised = out
return (out, out_denoised)
class LanPaint_SamplerCustomAdvanced:
@classmethod
def INPUT_TYPES(s):
return {"required":
{"noise": ("NOISE",),
"guider": ("GUIDER",),
"sampler": ("SAMPLER",),
"sigmas": ("SIGMAS",),
"latent_image": ("LATENT",),
"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
"return_with_leftover_noise": (["disable", "enable"], ),
"guider": ("GUIDER", ),
"sampler": ("SAMPLER", ),
"sigmas": ("SIGMAS", ),
"latent_image": ("LATENT", ),
"LanPaint_NumSteps": ("INT", {"default": 5, "min": 0, "max": 100, "tooltip": "Number of steps for Langevin dynamics, representing turns of thinking per step."}),
"LanPaint_Lambda": ("FLOAT", {"default": 16.0, "min": 0.1, "max": 50.0, "step": 0.1, "tooltip": "Bidirectional guidance scale. Higher values align with known regions but may cause instability."}),
"LanPaint_StepSize": ("FLOAT", {"default": 0.15, "min": 0.0001, "max": 1.0, "step": 0.01, "tooltip": "Step size for Langevin dynamics. Higher values speed convergence but may be unstable."}),
"LanPaint_StepSize": ("FLOAT", {"default": 0.2, "min": 0.0001, "max": 1.0, "step": 0.01, "tooltip": "Step size for Langevin dynamics. Higher values speed convergence but may be unstable."}),
"LanPaint_Beta": ("FLOAT", {"default": 1.0, "min": 0.0001, "max": 5.0, "step": 0.1, "tooltip": "Step size ratio between masked/unmasked regions. Lower values balance high Lambda."}),
"LanPaint_Friction": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 50.0, "step": 0.1, "tooltip": "Friction parameter for fast Langevin. Lower values speed convergence but may be unstable."}),
"LanPaint_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: prioritizes image quality; Prompt First: prioritizes prompt adherence."}),
"LanPaint_EarlyStop": ("INT", {"default": 1, "min": 0, "max": 10000, "tooltip": "Steps to stop LanPaint early, preventing irregular patterns."}),
"LanPaint_Info": ("STRING", {"default": "LanPaint Custom Sampler Adv. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
"LanPaint_InnerThreshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.0001, "round": 0.0001, "tooltip": "Early stop threshold for Langevin iterations based on semantic distance. 0.0 to disable. (Contributed by godnight10061)"}),
"LanPaint_InnerPatience": ("INT", {"default": 1, "min": 1, "max": 100, "tooltip": "Number of consecutive steps below threshold required to stop. (Contributed by godnight10061)"}),
}
}
RETURN_TYPES = ("LATENT", "LATENT")
RETURN_TYPES = ("LATENT","LATENT")
RETURN_NAMES = ("output", "denoised_output")
FUNCTION = "sample"
CATEGORY = "sampling/custom_sampling"
def sample(self, noise, guider, sampler, sigmas, latent_image, start_at_step, end_at_step, return_with_leftover_noise, LanPaint_NumSteps, LanPaint_Lambda, LanPaint_StepSize, LanPaint_Beta, LanPaint_Friction, LanPaint_PromptMode, LanPaint_EarlyStop, LanPaint_Info=""):
force_full_denoise = True
if end_at_step <= start_at_step:
raise ValueError('end_at_step must be larger than start_at_step')
if return_with_leftover_noise == "enable":
force_full_denoise = False
def sample(self, noise, guider, sampler, sigmas, latent_image, LanPaint_NumSteps, LanPaint_Lambda, LanPaint_StepSize, LanPaint_Beta, LanPaint_Friction, LanPaint_PromptMode, LanPaint_EarlyStop, LanPaint_Info="", LanPaint_InnerThreshold=0.0, LanPaint_InnerPatience=1):
model = guider.model_patcher
model.LanPaint_StepSize = LanPaint_StepSize
model.LanPaint_Lambda = LanPaint_Lambda
@@ -497,52 +592,40 @@ class LanPaint_SamplerCustomAdvanced:
model.LanPaint_NumSteps = LanPaint_NumSteps
model.LanPaint_Friction = LanPaint_Friction
model.LanPaint_EarlyStop = LanPaint_EarlyStop
model.LanPaint_InnerThreshold = LanPaint_InnerThreshold
model.LanPaint_InnerPatience = LanPaint_InnerPatience
if LanPaint_PromptMode == "Image First":
model.LanPaint_cfg_BIG = guider.cfg
else:
model.LanPaint_cfg_BIG = 0 * guider.cfg - 0.5
with override_sample_function():
latent = latent_image.copy()
latent_image_samples = latent["samples"]
latent_image_samples = comfy.sample.fix_empty_latent_channels(model, latent_image_samples)
latent["samples"] = latent_image_samples
latent = latent_image
latent_image = latent["samples"]
latent = latent.copy()
latent_image = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image)
latent["samples"] = latent_image
noise_mask = None
if "noise_mask" in latent:
noise_mask = latent["noise_mask"]
# From base comfy samplers.py
if end_at_step is not None and end_at_step < (len(sigmas) - 1):
sigmas = sigmas[:end_at_step + 1]
if force_full_denoise:
sigmas[-1] = 0
if start_at_step is not None:
if start_at_step < (len(sigmas) - 1):
sigmas = sigmas[start_at_step:]
else:
if latent_image is not None:
return latent_image
else:
return torch.zeros_like(noise)
x0_output = {}
callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output)
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
samples = guider.sample( noise.generate_noise(latent), latent_image_samples, sampler, sigmas, denoise_mask=noise_mask, callback=callback,disable_pbar=disable_pbar, seed=noise.seed
)
samples = guider.sample(noise.generate_noise(latent), latent_image, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed)
samples = samples.to(comfy.model_management.intermediate_device())
out = latent.copy()
out["samples"] = samples
if "x0" in x0_output:
out_denoised = latent.copy()
out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
out_denoised["samples"] = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
else:
out_denoised = out
return (out, out_denoised)
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
+10
View File
@@ -0,0 +1,10 @@
from typing import NamedTuple, Optional
import torch
class LangevinState(NamedTuple):
v: Optional[torch.Tensor]
C: Optional[torch.Tensor]
x0: Optional[torch.Tensor]
+20 -21
View File
@@ -28,11 +28,11 @@ def expm1mxmhx2_x3(x):
def exp_1mcosh_GD(gamma_t, delta):
"""
Compute e^(-Γt) * (1 - cosh(Γt√Δ))/ ( (Γt)**2 Δ )
Parameters:
gamma_t: Γ*t term (could be a scalar or tensor)
delta: Δ term (could be a scalar or tensor)
Returns:
Result of the computation with numerical stability handling
"""
@@ -68,7 +68,6 @@ def exp_sinh_GsqrtD(gamma_t, delta):
sqrt_abs_delta = torch.sqrt(torch.abs(delta))
gamma_t_sqrt_delta = gamma_t * sqrt_abs_delta
numerator_pos = (torch.exp(gamma_t * (sqrt_abs_delta - 1)) - torch.exp(gamma_t * (-sqrt_abs_delta - 1))) / 2
denominator_pos = gamma_t_sqrt_delta
result_pos = numerator_pos / gamma_t_sqrt_delta
result_pos = torch.where(torch.isfinite(result_pos), result_pos, torch.zeros_like(result_pos))
@@ -117,15 +116,15 @@ def zeta1(gamma_t, delta):
exp_cosh_term = exp_cosh(half_gamma_t, delta)
exp_sinh_term = exp_sinh_sqrtD(half_gamma_t, delta)
# Main computation
numerator = 1 - (exp_cosh_term + exp_sinh_term)
denominator = gamma_t * (1 - delta) / 4
result = 1 - numerator / denominator
# Handle numerical instability
result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
# Taylor expansion for small x (similar to your epxm1Dx approach)
mask = torch.abs(denominator) < 5e-3
term1 = epxm1_x(-gamma_t)
@@ -133,17 +132,17 @@ def zeta1(gamma_t, delta):
term3 = expm1mxmhx2_x3(-gamma_t)
taylor = term1 + (1/2.+ term1-3*term2)*denominator + (-1/6. + term1/2 - 4 * term2 + 10 * term3) * denominator**2
result = torch.where(mask, taylor, result)
return result
def exp_cosh_minus_terms(gamma_t, delta):
"""
Compute E^(-tΓ) * (Cosh[tΓ] - 1 - (Cosh[tΓ√Δ] - 1)/Δ) / (tΓ(1 - Δ))
Parameters:
gamma_t: Γ*t term (could be a scalar or tensor)
delta: Δ term (could be a scalar or tensor)
Returns:
Result of the computation with numerical stability handling
"""
@@ -151,17 +150,17 @@ def exp_cosh_minus_terms(gamma_t, delta):
# Compute individual terms
exp_cosh_term = exp_cosh(gamma_t, gamma_t**0) - exp_term # E^(-tΓ) (Cosh[tΓ] - 1) term
exp_cosh_delta_term = - gamma_t**2 * exp_1mcosh_GD(gamma_t, delta) # E^(-tΓ) (Cosh[tΓ√Δ] - 1)/Δ term
#exp_1mcosh_GD e^(-Γt) * (1 - cosh(Γt√Δ))/ ( (Γt)**2 Δ )
# Main computation
numerator = exp_cosh_term - exp_cosh_delta_term
denominator = gamma_t * (1 - delta)
result = numerator / denominator
# Handle numerical instability
result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
# Taylor expansion for small gamma_t and delta near 1
mask = (torch.abs(denominator) < 1e-1)
exp_1mcosh_GD_term = exp_1mcosh_GD(gamma_t, delta**0)
@@ -170,7 +169,7 @@ def exp_cosh_minus_terms(gamma_t, delta):
- denominator / 4 * ( 0.5 * exp_cosh(gamma_t, delta**0) - 4 * exp_1mcosh_GD_term - 5 /2 * exp_sinh_GsqrtD(gamma_t, delta**0) )
)
result = torch.where(mask, taylor, result)
return result
@@ -185,7 +184,7 @@ def sig11(gamma_t, delta):
def Zcoefs(gamma_t, delta):
Zeta1 = zeta1(gamma_t, delta)
Zeta2 = zeta2(gamma_t, delta)
sq_total = 1 - Zeta1 + gamma_t * (delta - 1) * (Zeta1 - 1)**2 / 8
amplitude = torch.sqrt(sq_total)
Zcoef1 = ( gamma_t**0.5 * Zeta2 / 2 **0.5 ) / amplitude
@@ -208,7 +207,7 @@ class StochasticHarmonicOscillator:
dq(t) = -Γ A y(t) dt + Γ C dt + Γ D dw(t) - Γ q(t) dt
Also define v(t) = q(t) / √Γ, which is numerically more stable.
Where:
y(t) - Position variable
q(t) - Velocity variable
@@ -239,7 +238,7 @@ class StochasticHarmonicOscillator:
Returns:
tuple: (y(t), v(t))
"""
dummyzero = y0.new_zeros(1) # convert scalar to tensor with same device and dtype as y0
Delta = self.Delta + dummyzero
Gamma_hat = self.Gamma * t + dummyzero
@@ -254,12 +253,12 @@ class StochasticHarmonicOscillator:
if v0 is None:
v0 = torch.randn_like(y0) * D / 2 ** 0.5
#v0 = (C - A * y0)/Gamma**0.5
# Calculate mean position and velocity
term1 = (1 - zeta_1) * (C * t - A * t * y0) + zeta_2 * (Gamma ** 0.5) * v0 * t
y_mean = term1 + y0
v_mean = (1 - EE)*(C - A * y0) / (Gamma ** 0.5) + (EE - A * t * (1 - zeta_1)) * v0
cov_yy = D**2 * t * self.sig22(Gamma_hat, Delta)
cov_vv = D**2 * self.sig11(Gamma_hat, Delta) / 2
cov_yv = (zeta2(Gamma_hat, Delta) * Gamma_hat * D ) **2 / 2 / (Gamma ** 0.5)
@@ -274,7 +273,7 @@ class StochasticHarmonicOscillator:
cov_matrix[..., 1, 1] = cov_vv
# Compute the Cholesky decomposition to get scale_tril
#scale_tril = torch.linalg.cholesky(cov_matrix)
scale_tril = torch.zeros(*batch_shape, 2, 2, device=y0.device, dtype=y0.dtype)
@@ -298,4 +297,4 @@ class StochasticHarmonicOscillator:
scale_tril=scale_tril
).sample()
return new_yv[...,0], new_yv[...,1]
return new_yv[...,0], new_yv[...,1]
+4 -3
View File
@@ -1,4 +1,5 @@
[pytest]
testpaths = . # Run tests in the current directory
python_files = test_*.py # Run tests in files that start with "test_"
norecursedirs = .. # Don't run tests in the parent directory
# Keep settings value-only; pytest does not treat inline `# ...` as comments.
testpaths = .
python_files = test_*.py
norecursedirs = ..
+9 -17
View File
@@ -1,21 +1,13 @@
#!/usr/bin/env python
"""Basic import tests for LanPaint.
"""Tests for `LanPaint` package."""
The ComfyUI runtime dependencies (e.g. `comfy`) are intentionally optional for unit tests.
"""
import pytest
from src.LanPaint.nodes import Example
@pytest.fixture
def example_node():
"""Fixture to create an Example node instance."""
return Example()
def test_package_imports_without_comfy() -> None:
import LanPaint
def test_example_node_initialization(example_node):
"""Test that the node can be instantiated."""
assert isinstance(example_node, Example)
def test_return_types():
"""Test the node's metadata."""
assert Example.RETURN_TYPES == ("IMAGE",)
assert Example.FUNCTION == "test"
assert Example.CATEGORY == "Example"
assert isinstance(LanPaint.NODE_CLASS_MAPPINGS, dict)
assert isinstance(LanPaint.NODE_DISPLAY_NAME_MAPPINGS, dict)
assert "LanPaint_KSampler" in LanPaint.NODE_CLASS_MAPPINGS
assert LanPaint.WEB_DIRECTORY == "./web"
+104
View File
@@ -0,0 +1,104 @@
import torch
from src.LanPaint.lanpaint import LanPaint as LanPaintEngine
class _DummySampling:
def noise_scaling(self, sigma, noise, latent_image): # type: ignore[no-untyped-def]
return latent_image + noise * sigma
class _DummyModel:
def __init__(self) -> None:
self.inner_model = self
self.model_sampling = _DummySampling()
def __call__(self, x, sigma, model_options=None, seed=None): # type: ignore[no-untyped-def]
return x, x
def _inputs(): # type: ignore[no-untyped-def]
x = torch.zeros((1, 4, 8, 8))
latent_image = torch.zeros_like(x)
noise = torch.ones_like(x)
sigma = torch.tensor([1.0])
latent_mask = torch.zeros_like(x)
current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
return x, latent_image, noise, sigma, latent_mask, current_times
def test_default_semantic_stop_triggers_at_patience_without_custom_distance_fn() -> None:
engine = LanPaintEngine(
_DummyModel(),
NSteps=10,
Friction=15.0,
Lambda=1.0,
Beta=1.0,
StepSize=0.2,
)
calls = {"langevin": 0, "with_score": 0, "without_score": 0}
def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None): # type: ignore[no-untyped-def]
calls["langevin"] += 1
if score is None:
calls["without_score"] += 1
else:
calls["with_score"] += 1
return x_t, args
engine.langevin_dynamics = fake_langevin # type: ignore[method-assign]
model_options = {
"lanpaint_semantic_stop": {
"threshold": 1e-6,
"patience": 2,
}
}
x, latent_image, noise, sigma, latent_mask, current_times = _inputs()
engine(x, latent_image, noise, sigma, latent_mask, current_times, model_options=model_options, seed=0, n_steps=10)
assert calls["langevin"] == 3
assert calls["with_score"] == 3
assert calls["without_score"] == 0
def test_semantic_stop_is_disabled_when_no_inpaint_region() -> None:
engine = LanPaintEngine(
_DummyModel(),
NSteps=10,
Friction=15.0,
Lambda=1.0,
Beta=1.0,
StepSize=0.2,
)
calls = {"langevin": 0, "with_score": 0, "without_score": 0}
def fake_langevin(x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None): # type: ignore[no-untyped-def]
calls["langevin"] += 1
if score is None:
calls["without_score"] += 1
else:
calls["with_score"] += 1
return x_t, args
engine.langevin_dynamics = fake_langevin # type: ignore[method-assign]
model_options = {
"lanpaint_semantic_stop": {
"threshold": 1e-6,
"patience": 1,
}
}
x, latent_image, noise, sigma, latent_mask, _ = _inputs()
current_times = (sigma, torch.tensor([0.5]), torch.tensor([0.0]))
no_inpaint_mask = torch.ones_like(latent_mask)
engine(x, latent_image, noise, sigma, no_inpaint_mask, current_times, model_options=model_options, seed=0, n_steps=10)
assert calls["langevin"] == 10
assert calls["with_score"] == 10
assert calls["without_score"] == 0
+74
View File
@@ -0,0 +1,74 @@
import importlib
import sys
import types
import pytest
import torch
def _repeat_to_batch_size(tensor: torch.Tensor, batch_size: int) -> torch.Tensor:
if tensor.shape[0] == batch_size:
return tensor
if tensor.shape[0] == 1:
return tensor.repeat((batch_size,) + (1,) * (tensor.ndim - 1))
repeats = (batch_size + tensor.shape[0] - 1) // tensor.shape[0]
return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1))[:batch_size]
def _import_nodes(monkeypatch, comfyui_version: str):
comfy_mod = types.ModuleType("comfy")
comfy_mod.__path__ = []
comfy_utils_mod = types.ModuleType("comfy.utils")
comfy_utils_mod.repeat_to_batch_size = _repeat_to_batch_size
comfy_samplers_mod = types.ModuleType("comfy.samplers")
class DummyKSAMPLER: ...
comfy_samplers_mod.KSAMPLER = DummyKSAMPLER
comfy_model_base_mod = types.ModuleType("comfy.model_base")
class ModelType:
FLUX = "FLUX"
FLOW = "FLOW"
class WAN22: ...
comfy_model_base_mod.ModelType = ModelType
comfy_model_base_mod.WAN22 = WAN22
comfyui_version_mod = types.ModuleType("comfyui_version")
comfyui_version_mod.__version__ = comfyui_version
comfy_mod.utils = comfy_utils_mod
comfy_mod.samplers = comfy_samplers_mod
comfy_mod.model_base = comfy_model_base_mod
monkeypatch.setitem(sys.modules, "comfy", comfy_mod)
monkeypatch.setitem(sys.modules, "comfy.utils", comfy_utils_mod)
monkeypatch.setitem(sys.modules, "comfy.samplers", comfy_samplers_mod)
monkeypatch.setitem(sys.modules, "comfy.model_base", comfy_model_base_mod)
monkeypatch.setitem(sys.modules, "nodes", types.ModuleType("nodes"))
monkeypatch.setitem(sys.modules, "latent_preview", types.ModuleType("latent_preview"))
monkeypatch.setitem(sys.modules, "comfyui_version", comfyui_version_mod)
sys.modules.pop("src.LanPaint.nodes", None)
return importlib.import_module("src.LanPaint.nodes")
@pytest.mark.parametrize("comfyui_version", ["0.5.0", "0.6.0"])
def test_reshape_mask_accepts_bhw_and_5d_output_shape(monkeypatch, comfyui_version: str) -> None:
lanpaint_nodes = _import_nodes(monkeypatch, comfyui_version)
input_mask = torch.zeros((1, 4, 4))
output_shape = (1, 16, 1, 8, 8)
out = lanpaint_nodes.reshape_mask(input_mask, output_shape, video_inpainting=False)
assert tuple(out.shape) == output_shape
def test_prepare_mask_accepts_hw_and_moves_device(monkeypatch) -> None:
lanpaint_nodes = _import_nodes(monkeypatch, "0.5.0")
input_mask = torch.zeros((4, 4))
output_shape = (2, 3, 8, 8)
out = lanpaint_nodes.prepare_mask(input_mask, output_shape, device=torch.device("cpu"), video_inpainting=False)
assert tuple(out.shape) == output_shape
assert out.device.type == "cpu"
+45
View File
@@ -0,0 +1,45 @@
import torch
from unittest.mock import MagicMock, patch
from src.LanPaint.lanpaint import LanPaint
def test_langevin_dynamics_fallback_on_nan() -> None:
"""Test that langevin_dynamics falls back to overdamped dynamics if damped dynamics produces NaNs."""
torch.manual_seed(0)
# Setup minimal LanPaint instance
lp = LanPaint(Model=MagicMock(), NSteps=10, Friction=1.0, Lambda=1.0, Beta=1.0, StepSize=0.1)
# Dummy inputs
# Shape: (Batch, Channel, Height, Width)
x_t = torch.randn(1, 4, 8, 8)
lp.img_dim_size = 4
mask = torch.zeros_like(x_t)
# Simple score function
def score(x):
return torch.zeros_like(x)
step_size = torch.tensor([0.1])
# (sigma, abt, flow_t)
current_times = (torch.tensor([0.5]), torch.tensor([0.5]), torch.tensor([0.5]))
# Mock StochasticHarmonicOscillator to return NaNs
# We patch it where it is used (imported) in lanpaint.py
with patch("src.LanPaint.lanpaint.StochasticHarmonicOscillator") as MockSHO:
mock_instance = MockSHO.return_value
# Configure dynamics to return NaNs
nan_tensor = torch.full_like(x_t, float('nan'))
mock_instance.dynamics.return_value = (nan_tensor, nan_tensor)
# Execute langevin_dynamics
# This should try run_damped -> get NaNs -> raise ValueError -> catch -> run_overdamped
x_out, args_out = lp.langevin_dynamics(x_t, score, mask, step_size, current_times, sigma_y=1.0)
assert hasattr(args_out, "v")
assert hasattr(args_out, "C")
assert hasattr(args_out, "x0")
assert args_out[0] is args_out.v
assert args_out[1] is args_out.C
assert args_out[2] is args_out.x0
v_out = args_out[0]
# Verify that SHO was initialized and dynamics called
MockSHO.assert_called()
mock_instance.dynamics.assert_called()
# Verify result is finite (indicating fallback to overdamped logic was successful)
assert torch.isfinite(x_out).all(), "Output contains NaNs, fallback failed"
assert v_out is None or torch.isfinite(v_out).all()