Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
75e606579c | ||
|
|
6440a7f433 | ||
|
|
b3dc7e7551 | ||
|
|
2522687dd5 | ||
|
|
f3ef9dedf9 | ||
|
|
f53442c8d5 | ||
|
|
24051fb151 | ||
|
|
a76a5da0c5 | ||
|
|
21a6b155e8 | ||
|
|
8bc20d5e00 | ||
|
|
84b63adb11 | ||
|
|
725fb2ba37 | ||
|
|
4f19674299 | ||
|
|
558084416d | ||
|
|
a984dcb118 | ||
|
|
3f9193997a | ||
|
|
7372428114 | ||
|
|
ac6eece747 | ||
|
|
95b65ed1b1 | ||
|
|
30cfd2f2e4 | ||
|
|
2e3a566bb4 | ||
|
|
8fe1a9d312 | ||
|
|
4fcbfc7b5b | ||
|
|
651ddc2225 | ||
|
|
bd4606f27a | ||
|
|
8a72294f27 | ||
|
|
9ddd5f613e | ||
|
|
39c90bbf62 | ||
|
|
6ad4a18537 | ||
|
|
5e7d472738 | ||
|
|
d73f3db265 | ||
|
|
cd06069b52 | ||
|
|
114d091965 | ||
|
|
06da285a09 | ||
|
|
8f52a1d378 | ||
|
|
c7962bf784 | ||
|
|
a895d2a758 | ||
|
|
34ded92d8d | ||
|
|
1b87239553 | ||
|
|
cbc9f3abfd | ||
|
|
7d46147887 | ||
|
|
c5fe6351ea | ||
|
|
ac7a5a1f2f | ||
|
|
73ab537e86 | ||
|
|
6fd5f60383 | ||
|
|
9c9b5555fc | ||
|
|
450a25e5e8 | ||
|
|
e7080f9574 | ||
|
|
2086f00603 | ||
|
|
896729ccd4 | ||
|
|
207860edc3 | ||
|
|
4074473b90 | ||
|
|
11c1490140 | ||
|
|
7c79712fc8 | ||
|
|
fb35a1d032 | ||
|
|
060fcc4475 | ||
|
|
aedb908a3f | ||
|
|
f9718ea3e1 | ||
|
|
0bafe2117a | ||
|
|
2c1a23777d | ||
|
|
dde4c82463 | ||
|
|
72cb484c40 | ||
|
|
0d15d15cf3 | ||
|
|
1393a46d67 | ||
|
|
27421363bf | ||
|
|
01400a541d | ||
|
|
c981354387 | ||
|
|
c92c8bfb1f | ||
|
|
9334025801 | ||
|
|
144893ecd2 | ||
|
|
bce8505ce2 | ||
|
|
0992657e54 | ||
|
|
2fa036992b | ||
|
|
9432a34c38 | ||
|
|
d4e8ee28fb | ||
|
|
34b8e83831 | ||
|
|
e29d480a1a | ||
|
|
9479106c1d | ||
|
|
db4bbbe9f2 | ||
|
|
1a3e15964c | ||
|
|
4cc30f8a1d | ||
|
|
d1e609192f | ||
|
|
bf46831196 | ||
|
|
c9d9a79b18 | ||
|
|
175156af86 | ||
|
|
3281946b32 | ||
|
|
83f557ff4c | ||
|
|
5e7fe4d5e4 | ||
|
|
69f2aded4d | ||
|
|
5f1b6d8989 | ||
|
|
7b0d144db9 | ||
|
|
27ebb6e7af | ||
|
|
f148e4b631 | ||
|
|
cc9ec8873a | ||
|
|
565087e8f8 | ||
|
|
ebdc5dcc02 | ||
|
|
9336523576 | ||
|
|
21520b38fb | ||
|
|
868bbaeefe | ||
|
|
3581d6c217 | ||
|
|
d4da280a92 | ||
|
|
d0bcd459cc | ||
|
|
469012d183 | ||
|
|
3f4a5cdb5a | ||
|
|
e0ce52b1be | ||
|
|
2c8f4baec0 | ||
|
|
daf1bdf4be | ||
|
|
f4d4dfdd61 | ||
|
|
6d624e7e1e | ||
|
|
e0940059e6 | ||
|
|
76ce20b51f | ||
|
|
29ed0b3da5 | ||
|
|
3d5aa3a3fa | ||
|
|
6109df6591 | ||
|
|
536716a7b9 | ||
|
|
ece228ab10 | ||
|
|
c048230250 | ||
|
|
aa885d9fa6 | ||
|
|
80d55a47b0 | ||
|
|
79db1765e3 | ||
|
|
929328d08d | ||
|
|
c12897040b | ||
|
|
48b9a59ca3 | ||
|
|
d55f8c61d0 | ||
|
|
a097b9c9b2 | ||
|
|
70cc94cc6a | ||
|
|
e653f45103 | ||
|
|
c9c237c2de | ||
|
|
c5aa0116d5 |
@@ -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 .
|
||||
|
||||
@@ -11,3 +11,5 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: comfy-org/node-diff@main
|
||||
with:
|
||||
base_ref: ${{ github.event.repository.default_branch }}
|
||||
|
||||
@@ -98,3 +98,5 @@ cookiecutter-pypackage-env/
|
||||
# vscode settings
|
||||
.history/
|
||||
*.code-workspace
|
||||
.vscode/
|
||||
/.vscode
|
||||
|
||||
@@ -5,4 +5,8 @@
|
||||
"/PATH/TO/ComfyUI/",
|
||||
"/PATH/TO/ComfyUI/custom_nodes/"
|
||||
],
|
||||
"cursorpyright.analysis.extraPaths": [
|
||||
"/PATH/TO/ComfyUI/",
|
||||
"/PATH/TO/ComfyUI/custom_nodes/"
|
||||
],
|
||||
}
|
||||
|
||||
@@ -1,49 +1,115 @@
|
||||
<div align="center">
|
||||
|
||||
# LanPaint: Universal Inpainting Sampler with "Think Mode"
|
||||
[](https://arxiv.org/abs/2502.03491)
|
||||
[](https://openreview.net/pdf?id=JPC8JyOUSW)
|
||||
[](https://github.com/scraed/LanPaintBench)
|
||||
[](https://github.com/comfyanonymous/ComfyUI)
|
||||
[](https://huggingface.co/charrywhite/LanPaint)
|
||||
[](https://scraed.github.io/scraedBlog/)
|
||||
[](https://github.com/scraed/LanPaint/stargazers)
|
||||
[](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.
|
||||
|
||||

|
||||
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.
|
||||
|
||||

|
||||
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 |
|
||||
|:--------:|:------:|:---------:|
|
||||
|  |  |  |
|
||||
|
||||
**🎬 NEW: LanPaint now supports Z-Image-Base too!**
|
||||
|
||||
| Original | Masked | Inpainted |
|
||||
|:--------:|:------:|:---------:|
|
||||
|  |  |  |
|
||||
|
||||
|
||||
**🎬 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 |
|
||||
|:--------------:|:----:|:----------------:|
|
||||
|  |  |  |
|
||||
|
||||
*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.
|
||||

|
||||
- **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 |
|
||||
|:--------------:|:----:|:----------------:|
|
||||
|  |  |  |
|
||||
|
||||
[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 |
|
||||
|:--------------:|:----:|:-----------------:|
|
||||
|  |  |  |
|
||||
|
||||
[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 |
|
||||
|:--------:|:------:|:---------:|
|
||||
|  |  |  |
|
||||
|
||||
</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 |
|
||||
|:--------:|:------:|:----------:|
|
||||
|  |  |  |
|
||||
|
||||
</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 |
|
||||
|:--------:|:------:|:---------:|
|
||||
|  |  |  |
|
||||
|
||||
</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.
|
||||
|
||||

|
||||
[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.
|
||||
|
||||

|
||||
|
||||
### Example Qwen Edit 2508: InPaint
|
||||

|
||||
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.)
|
||||
|
||||

|
||||
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)
|
||||

|
||||
[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 |
|
||||
|:--------:|:------:|:---------:|
|
||||
|  |  |  |
|
||||
|
||||
</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 |
|
||||
|:--------:|:------:|:---------:|
|
||||
|  |  |  |
|
||||
|
||||
</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)
|
||||

|
||||
[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
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
After Width: | Height: | Size: 1.6 MiB |
|
After Width: | Height: | Size: 1.6 MiB |
|
After Width: | Height: | Size: 1.3 MiB |
@@ -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
|
||||
}
|
||||
|
After Width: | Height: | Size: 670 KiB |
|
After Width: | Height: | Size: 674 KiB |
|
After Width: | Height: | Size: 1.1 MiB |
|
After Width: | Height: | Size: 1.1 MiB |
|
After Width: | Height: | Size: 687 KiB |
|
After Width: | Height: | Size: 1.4 MiB |
|
After Width: | Height: | Size: 1.4 MiB |
|
After Width: | Height: | Size: 1.6 MiB |
|
After Width: | Height: | Size: 1.6 MiB |
|
After Width: | Height: | Size: 221 KiB |
|
After Width: | Height: | Size: 141 KiB |
|
After Width: | Height: | Size: 1.3 MiB |
|
After Width: | Height: | Size: 1.6 MiB |
|
After Width: | Height: | Size: 1.6 MiB |
|
After Width: | Height: | Size: 1.1 MiB |
|
After Width: | Height: | Size: 1.3 MiB |
|
After Width: | Height: | Size: 1.4 MiB |
|
After Width: | Height: | Size: 1.9 MiB |
|
After Width: | Height: | Size: 1.2 MiB |
|
After Width: | Height: | Size: 1.4 MiB |
|
After Width: | Height: | Size: 1.6 MiB |
|
After Width: | Height: | Size: 1.8 MiB |
|
After Width: | Height: | Size: 1.6 MiB |
|
After Width: | Height: | Size: 1.8 MiB |
|
After Width: | Height: | Size: 1.9 MiB |
|
After Width: | Height: | Size: 1.7 MiB |
|
After Width: | Height: | Size: 2.2 MiB |
|
After Width: | Height: | Size: 1.8 MiB |
|
After Width: | Height: | Size: 2.8 MiB |
|
After Width: | Height: | Size: 461 KiB |
|
After Width: | Height: | Size: 4.1 MiB |
|
After Width: | Height: | Size: 5.8 MiB |
|
After Width: | Height: | Size: 801 KiB |
|
After Width: | Height: | Size: 551 KiB |
|
After Width: | Height: | Size: 6.2 MiB |
|
After Width: | Height: | Size: 4.8 MiB |
|
After Width: | Height: | Size: 4.1 MiB |
|
After Width: | Height: | Size: 8.5 MiB |
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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 = ..
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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()
|
||||
|
||||