Compare commits

..
128 Commits
Author SHA1 Message Date
scraed 7aeb6e535f Add Qwen image support and new inpainting examples
Updated the README to announce Qwen image support and provide usage instructions. Added new example images and workflow files for Qwen inpainting in the examples directory.
2025-08-08 14:09:53 +08:00
scraed eda0f19944 Generalize dimension handling in LanPaint and fix mask reshape
Refactored LanPaint to dynamically handle tensors with varying numbers of dimensions by introducing add_none_dims and remove_none_dims utility methods. Updated all relevant tensor broadcasting to use these methods, improving flexibility. Also fixed reshape_mask in nodes.py to use the last two dimensions for resizing, ensuring correct mask shape.
2025-08-08 13:42:24 +08:00
scraed e20c8f20ce update version 2025-06-21 14:13:09 +08:00
scraed 89b9010b35 Fix indentation error in sample method
Corrected an indentation issue in the sample method of LanPaint_SamplerCustomAdvanced to ensure proper execution of the end_at_step check.
2025-06-21 14:09:11 +08:00
scraed ee0c65656e simplify custom sampler 2025-06-21 13:18:25 +08:00
scraed 48b1dd4be6 Merge branch 'master' into pr/33 2025-06-21 13:10:05 +08:00
scraed 9d304cd5a3 update readme 2025-06-21 13:03:30 +08:00
scraed 61c19ac31d update algorithm with better outpaint 2025-06-21 12:56:36 +08:00
Bagier 1ca6e53090 adjust info names 2025-06-20 08:31:51 +07:00
Bagier 520932cee0 add invalid value check 2025-06-19 20:49:20 +07:00
Bagier 37612bb399 Logic fix 2025-06-19 20:12:50 +07:00
Bagier 0e6cb45081 re-add lanpaint info 2025-06-19 19:25:59 +07:00
Bagier 0a5e36d55e add advanced custom sampler 2025-06-19 19:15:58 +07:00
Bagier b319f772cb fix parameter update error 2025-06-19 19:00:41 +07:00
Bagier 340376ad18 Add custom sampler variation 2025-06-19 16:14:26 +07:00
scraed dbfc1585fc Merge branch 'master' of https://github.com/scraed/LanPaint 2025-06-16 20:03:31 +08:00
scraed c3ae2c644d update coefficients 2025-06-16 20:03:25 +08:00
scraed c7017373c9 Update pyproject.toml 2025-06-08 17:59:04 +08:00
scraed 850f707eb6 Remove redundant comment 2025-06-08 17:58:40 +08:00
scraed 62870f060a Update README.md 2025-06-06 13:05:50 +08:00
scraed a91cefacf0 Update README.md 2025-06-06 13:05:16 +08:00
scraed 4265f71a85 Update README.md 2025-06-05 18:15:21 +08:00
scraed 4189ba80c2 add error message 2025-06-05 17:10:07 +08:00
scraed 23ad6e47fd add new masked blend node 2025-06-05 16:21:46 +08:00
scraed 4d3d5d17f0 Update README.md 2025-06-05 01:22:09 +08:00
scraed b86f3b7112 update version 2025-06-05 01:14:57 +08:00
scraed 96b7f3eecd fix sigma batch error 2025-06-05 01:11:04 +08:00
scraed e472574783 update readme 2025-06-04 21:33:58 +08:00
scraed 0b799be8ec support more sampler 2025-06-04 21:28:51 +08:00
scraed 82da0b4142 Merge branch 'master' of https://github.com/scraed/LanPaint 2025-06-04 20:04:58 +08:00
scraed c66155e783 add early stop 2025-06-04 20:04:53 +08:00
scraed d802f2f2b0 Update README.md fix typo 2025-06-04 14:34:10 +08:00
scraed 6153a450c9 Update README.md 2025-06-04 14:33:07 +08:00
scraed 493658d23a Update README.md 2025-06-04 14:28:11 +08:00
scraed 7de1054b39 Update README.md 2025-06-04 09:06:39 +08:00
scraed 574d9906ec Update README.md 2025-06-03 20:54:12 +08:00
scraed 61d4ecef75 Update README.md 2025-06-03 18:41:24 +08:00
scraed 5eb2d8ef48 Update README.md 2025-06-03 17:19:52 +08:00
scraed b668bcc42c Update README.md 2025-06-03 13:11:47 +08:00
scraed 05386f98d3 replace outdated flux picture 2025-06-03 11:49:40 +08:00
scraed ad5704d25f correct tricks for consistency 2025-06-03 11:06:35 +08:00
scraed 2ab72851b9 update summary img 2025-06-03 10:29:54 +08:00
scraed 14f7907c70 update version 2025-06-03 10:19:41 +08:00
scraed 5adce4ecdd update readme 2025-05-28 18:11:30 +08:00
scraed 2fd945e69b update readme and pictures 2025-05-28 18:04:52 +08:00
scraed 1bd6932cae change to ve notation and update examples 2025-05-28 17:22:15 +08:00
scraed 59b31303c8 update examples 2025-05-27 10:44:15 +08:00
scraed f906bd9cc9 reduce parameters 2025-05-22 15:01:36 +08:00
scraed 31e4909438 switch to sampler alg 2025-05-22 10:14:58 +08:00
scraed 6775e6ac37 1-a schedule and shared A 2025-05-17 19:06:43 +08:00
scraed 49b35bcb28 fix Zcoef asymp bug 2025-05-14 15:55:18 +08:00
scraed 13bd3182bf switch to separate file 2025-05-13 19:14:46 +08:00
scraed bb3e0078d4 switch to new new alg 2025-05-13 18:33:08 +08:00
scraed 5fc4cf2092 remove y time truncate 2025-05-12 22:34:44 +08:00
scraed 3f8e2b833c change y update alg 2025-05-12 22:27:56 +08:00
scraed 0f546f09b6 add truncate time step 2025-05-12 16:38:35 +08:00
scraed 99ac3ee06d Create utils.py 2025-05-12 11:40:05 +08:00
scraed 505fbfb7fc fix tamed bug 2025-05-12 09:49:12 +08:00
scraed aa541edaea change default lambda schedule to const 2025-05-03 23:16:07 +08:00
scraed 5382fbba21 remove redundant time step schedule 2025-05-03 22:31:59 +08:00
scraed 82a6798e29 change parameter range 2025-05-03 22:28:25 +08:00
scraed 3b7c79f431 add lamb schedule 2025-05-03 01:01:02 +08:00
scraed ae31ac7e70 complete migration to new alg formula 2025-05-02 22:31:37 +08:00
scraed 2a67dd353f update epxm1Dx 2025-05-02 21:14:54 +08:00
scraed 2d0f458695 update eps 2025-05-02 13:50:19 +08:00
scraed 6f2deeda51 separate ld 2025-05-02 00:38:58 +08:00
scraed 76cd0a5c1d sort tamed 2025-05-01 21:08:45 +08:00
scraed f90636d4c0 create hidream img 2025-04-23 22:05:02 +08:00
scraed dc4ef3aed4 Update README.md 2025-04-22 08:55:51 +08:00
scraed eca581ee79 Update README.md 2025-04-17 15:38:32 +08:00
scraed 75a91710f2 Update pyproject.toml 2025-04-16 23:50:17 +08:00
scraed f8eea61264 Update README.md 2025-04-16 23:48:58 +08:00
scraed e42d3b0426 Update README.md 2025-04-16 23:32:46 +08:00
scraed ee567b14e8 Update Example for Hidream 2025-04-16 23:21:55 +08:00
scraed 10a790974c Merge branch 'master' of https://github.com/scraed/LanPaint 2025-04-16 23:12:18 +08:00
scraed 00d2eb9885 Update hidream example 2025-04-16 23:12:13 +08:00
scraed b1139f285a Update Citationn 2025-04-16 22:58:40 +08:00
scraed eb69a21cf2 Update README.md for hidream support 2025-04-16 22:26:05 +08:00
scraed c863ade839 update version 2025-04-16 22:25:05 +08:00
scraed b07ed256c7 Primary support for Hidream 2025-04-16 22:24:18 +08:00
scraed 612cfb51f2 Merge pull request #19 from ComfyNodePRs/update-publish-yaml
Update Github Action for Publishing to Comfy Registry
2025-04-11 20:42:47 +08:00
scraed aa9319f7cf Update README.md 2025-03-22 21:26:08 +08:00
scraed a92d19d4f8 Updates dates 2025-03-22 21:23:12 +08:00
scraed 478b077c69 Update pyproject.toml 2025-03-22 21:18:12 +08:00
scraed 3015d6f0ef Update README.md 2025-03-22 21:17:17 +08:00
scraed 55cd84be93 Update README.md for flux and tease mode 2025-03-22 21:16:40 +08:00
scraed 758eae62c3 add flux example 2025-03-22 21:08:23 +08:00
scraed 58f8d01acf Merge branch 'master' of https://github.com/scraed/LanPaint 2025-03-22 20:14:27 +08:00
scraed 1684b616f1 Merge pull request #16 from charrywhite/master
Add preview hood in KSamplerX0Inpaint
2025-03-22 20:10:07 +08:00
scraed 65f98f7c48 disable CFG BIG for Flux 2025-03-22 20:05:30 +08:00
charrywhite 5d8e40a669 Add preview hood in KSamplerX0Inpaint 2025-03-22 19:57:59 +08:00
scraed ff1fb62d1a Merge branch 'master' of https://github.com/scraed/LanPaint 2025-03-22 19:46:42 +08:00
scraed b7d37ef121 flux support and parameter change
add flux support and eliminate parameters a and b
2025-03-22 19:46:40 +08:00
scraed 29c9d877be Update ToDo 2025-03-19 23:10:21 +08:00
scraed 6c1375c35f update control net support in readme 2025-03-13 23:10:49 +08:00
scraed 8c04ba410b More intuitive how it works 2025-03-13 22:29:46 +08:00
scraed a5c08aa0a4 Update how it works 2025-03-13 17:45:32 +08:00
scraed cefc2c43d3 add version requirement in README.md 2025-03-13 10:15:36 +08:00
scraed 62c5934a1a Update README.md 2025-03-12 17:10:47 +08:00
scraed 9ae04abb14 Update README.md 2025-03-12 10:15:22 +08:00
scraed 8446a1a735 Update README.md 2025-03-12 09:55:46 +08:00
scraed 439cc8371a Update tuning guide 2025-03-12 08:58:42 +08:00
scraed 0b6adec240 Update README.md 2025-03-11 23:46:52 +08:00
scraed 32dd4f23ff Update README.md 2025-03-11 22:49:04 +08:00
scraed ef927e6a9a Update README.md 2025-03-11 12:14:01 +08:00
scraed 142595b8ac Update README.md 2025-03-11 12:13:20 +08:00
scraed 26dd8dcf8e Update README.md 2025-03-11 12:08:09 +08:00
scraed c4e8ca9fb9 Update README.md 2025-03-11 10:54:51 +08:00
scraed 6073fb9846 Update README.md 2025-03-11 10:53:10 +08:00
scraed de01bc07d7 Update README.md 2025-03-11 10:49:02 +08:00
scraed 492114f14c Update README.md 2025-03-11 09:39:26 +08:00
charrywhite 47f4fea033 Update README.md 2025-03-11 09:24:02 +08:00
scraed 3dd2c28db3 Update README.md 2025-03-11 01:43:51 +08:00
scraed f58cfeb0ca Update README.md 2025-03-11 01:39:21 +08:00
scraed df43f9f078 Update README.md 2025-03-11 01:13:51 +08:00
scraed 3eebbf9ffd Update pyproject.toml 2025-03-11 01:09:45 +08:00
scraed 3fcf3e11ac Update nodes.py 2025-03-11 01:08:40 +08:00
scraed a9aff269a7 Update README.md 2025-03-11 01:06:16 +08:00
scraed 22cce34f79 Merge branch 'master' of https://github.com/scraed/LanPaint 2025-03-11 01:00:36 +08:00
scraed 95e566ca23 update nodes image 2025-03-11 01:00:09 +08:00
scraed bd318e4666 Update README.md 2025-03-11 00:52:47 +08:00
scraed 733cbc704f Update README.md 2025-03-11 00:49:56 +08:00
charrywhite 96ef5e4926 Update README.md
add more description for lanpaint
2025-03-11 00:46:04 +08:00
scraed bf554c1170 Merge branch 'master' of https://github.com/scraed/LanPaint 2025-03-11 00:45:26 +08:00
scraed 3dc4adeade update example 6 picture 2025-03-11 00:45:06 +08:00
scraed dd41680dd7 Update README.md 2025-03-11 00:40:21 +08:00
scraed 3fb85ab403 major update node functions and examples 2025-03-11 00:36:46 +08:00
snomiao 34d36e3c93 chore(publish): update workflow for node publishing
- Add permissions for issue writing in the workflow
- Set condition to run job only for 'scraed' repository owner
- Update action version from 'main' to 'v1' for stability and consistency
2025-03-04 04:12:56 +00:00
42 changed files with 1019 additions and 321 deletions
+5 -1
View File
@@ -7,14 +7,18 @@ on:
paths:
- "pyproject.toml"
permissions:
issues: write
jobs:
publish-node:
name: Publish Custom Node to registry
runs-on: ubuntu-latest
if: ${{ github.repository_owner == 'scraed' }}
steps:
- name: Check out code
uses: actions/checkout@v4
- name: Publish Custom Node
uses: Comfy-Org/publish-node-action@main
uses: Comfy-Org/publish-node-action@v1
with:
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} ## Add your own personal access token to your Github Repository secrets and reference it here.
BIN
View File
Binary file not shown.

Before

Width:  |  Height:  |  Size: 263 KiB

After

Width:  |  Height:  |  Size: 244 KiB

BIN
View File
Binary file not shown.

Before

Width:  |  Height:  |  Size: 178 KiB

After

Width:  |  Height:  |  Size: 119 KiB

+186 -63
View File
@@ -1,55 +1,35 @@
# LanPaint (Thinking mode Inpaint)
Unlock precise inpainting without additional training. LanPaint lets the model "think" through multiple iterations before denoising, aiming for seamless and accurate results. We encourage you to try it out and share your feedback through issues or discussions, as your input will help us enhance the algorithm's performance and stability.
Unlock precise inpainting without additional training. LanPaint lets the model "think" through multiple iterations before denoising, enabling you to invest more computation time for superior quality.
![Inpainting Result 13](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_13.jpg)
This is the official implementation of ["Lanpaint: Training-Free Diffusion Inpainting with Exact and Fast Conditional Inference"](https://arxiv.org/abs/2502.03491).
## Features
- 🎨 **Zero-Training Inpainting** - Works immediately with ANY SD model, even custom models you've trained yourself
- 🛠️ **Simple Integration** - Same workflow as standard ComfyUI KSampler
- 🚀 **Quality Enhancements** - High quality and seamless inpainting
- **Universal Compatibility** – Works instantly with almost any model (SD 1.5, XL, 3.5, Flux, HiDream, 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.
- **Flexible Masking** – Supports any mask shape, size, or position for inpainting/outpainting.
- **No Workarounds** – Generates 100% new content (no blending or smoothing) without relying on partial denoising.
- **Beyond Inpainting** – You can even use it as a simple way to generate consistent characters.
## Example Results
All examples use random seed 0 to ensure fair comparison.
### Example 1: Basket to Basket Ball (LanPaint K Sampler, It is fast).
![Inpainting Result 1](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_04.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_1)
[Model Used in This Example](https://civitai.com/models/1188071?modelVersionId=1408658)
### Example 2: White Shirt to Blue Shirt (LanPaint K Sampler (Advanced), it is slower.)
![Inpainting Result 2](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_05.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_2)
[Model Used in This Example](https://civitai.com/models/1188071?modelVersionId=1408658)
### Example 3: Smile to Sad (LanPaint K Sampler (Advanced))
![Inpainting Result 3](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_06.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_3)
[Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl)
### Example 4: Damage Restoration (LanPaint K Sampler (Advanced))
![Inpainting Result 3](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_07.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_4)
[Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl)
### Example 5: Huge Damage Restoration (LanPaint K Sampler (Advanced))
![Inpainting Result 3](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_08.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_5)
[Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl) (The key is increase LanPaint_stepsize to 0.5)
**Warning**: LanPaint has degraded performance on distillation models, such as Flux.dev, due to a similar [issue with LORA training](https://medium.com/@zhiwangshi28/why-flux-lora-so-hard-to-train-and-how-to-overcome-it-a0c70bc59eaf). Please use low flux guidance (1.0-2.0) to mitigate this [issue](https://github.com/scraed/LanPaint/issues/30).
**How to Use These Examples:**
1. Navigate to the **example** folder (i.e example_1) by clicking **View Workflow & Masks**, download all pictures.
2. Drag **InPainted_Drag_Me_to_ComfyUI.png** into ComfyUI to load the workflow.
3. Download the required model from Civitai by clicking **Model Used in This Example**.
4. Load the model into the **"Load Checkpoint"** node.
5. Upload **Original_No_Mask.png** to the **"Load image"** node in the **"Original Image"** group (far left).
6. Upload **Masked_Load_Me_in_Loader.png** to the **"Load image"** node in the **"Mask image for inpainting"** group (second from left).
7. Queue the task, you will get inpainted results from three methods:
- **[VAE Encode for Inpainting](https://comfyanonymous.github.io/ComfyUI_examples/inpaint/)** (middle),
- **[Set Latent Noise Mask](https://comfyui-wiki.com/en/tutorial/basic/how-to-inpaint-an-image-in-comfyui)** (second from right),
- **LanPaint** (far right).
## How It Works
LanPaint uses Langevin Dynamics as "thinking" steps, which digs deeper into the diffusion process and allows the model to generate more consistent results.
Compare and explore the results from each method!
LanPaint introduces "BIG score" that creates a **two-way alignment** between masked and unmasked areas. It continuously evaluates:
- *"Does the new content make sense with the existing elements?"*
- *"Do the existing elements support the new creation?"*
Based on this evaluation, LanPaint iteratively updates the noise in both the masked and unmasked regions.
LanPaint also implements an accurate, robust, and fast Langevin dynamics solver.
![WorkFlow](https://github.com/scraed/LanPaint/blob/master/Example.JPG)
## Quickstart
1. **Install ComfyUI**: Follow the official [ComfyUI installation guide](https://docs.comfy.org/get_started) to set up ComfyUI on your system.
1. **Install ComfyUI**: Follow the official [ComfyUI installation guide](https://docs.comfy.org/get_started) to set up ComfyUI on your system. Or ensure your ComfyUI version > 0.3.11.
2. **Install ComfyUI-Manager**: Add the [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) for easy extension management.
3. **Install LanPaint Nodes**:
- **Via ComfyUI-Manager**: Search for "[LanPaint](https://registry.comfy.org/publishers/scraed/nodes/LanPaint)" in the manager and install it directly.
@@ -62,6 +42,120 @@ Compare and explore the results from each method!
Once installed, you'll find the LanPaint nodes under the "sampling" category in ComfyUI. Use them just like the default KSampler for high-quality inpainting!
## Updates
- 2025/08/08
- Add Qwen image support
- 2025/06/21
- Update the algorithm with enhanced stability and outpaint performance.
- Add outpaint example
- Supports Sampler Custom (Thanks to [MINENEMA](https://github.com/MINENEMA))
- 2025/06/04
- Add more sampler support.
- Add early stopping to advanced sampler.
- 2025/05/28
- Major update on the Langevin solver. It is now much faster and more stable.
- Greatly simplified the parameters for advanced sampler.
- Fix performance issue on Flux and SD 3.5
- 2025/04/16
- Added Primary HiDream support
- 2025/03/22
- Added Primary Flux support
- Added Tease Mode
- 2025/03/10
- LanPaint has received a major update! All examples now use the LanPaint K Sampler, offering a simplified interface with enhanced performance and stability.
## Examples
### Example Qwen Image: InPaint(LanPaint K Sampler, 5 steps of thinking)
We are excited to announce that LanPaint now supports Qwen Image, providing powerful inpainting capabilities for image editing.
![Inpainting Result 14](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_14.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_11)
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.
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.)
### Example HiDream: InPaint (LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 8](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_11.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_8)
You need to follow the ComfyUI version of [HiDream workflow](https://docs.comfy.org/tutorials/image/hidream/hidream-i1) to download and install the model.
### Example HiDream: OutPaint(LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 8](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_13(1).jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_10)
You need to follow the ComfyUI version of [HiDream workflow](https://docs.comfy.org/tutorials/image/hidream/hidream-i1) to download and install the model. Thanks [Amazon90](https://github.com/Amazon90) for providing this example.
### Example SD 3.5: InPaint(LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 8](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_12.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_9)
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: InPaint(LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 7](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_10.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_7)
[Model Used in This Example](https://huggingface.co/Comfy-Org/flux1-dev/blob/main/flux1-dev-fp8.safetensors)
(Note: Prompt First mode is disabled on Flux. As it does not use CFG guidance.)
### Example SDXL 0: Character Consistency (Side View Generation) (LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 6](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_09.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_6)
[Model Used in This Example](https://civitai.com/models/1188071?modelVersionId=1408658)
(Tricks 1: You can emphasize the character by copy it's image multiple times with Photoshop. Here I have made one extra copy.)
(Tricks 2: Use prompts like multiple views, multiple angles, clone, turnaround. Use LanPaint's Prompt first mode (does not support Flux))
(Tricks 3: Remeber LanPaint can in-paint: Mask non-consistent regions and try again!)
### Example SDXL 1: Basket to Basket Ball (LanPaint K Sampler, 2 steps of thinking).
![Inpainting Result 1](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_04.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_1)
[Model Used in This Example](https://civitai.com/models/1188071?modelVersionId=1408658)
### Example SDXL 2: White Shirt to Blue Shirt (LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 2](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_05.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_2)
[Model Used in This Example](https://civitai.com/models/1188071?modelVersionId=1408658)
### Example SDXL 3: Smile to Sad (LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 3](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_06.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_3)
[Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl)
### Example SDXL 4: Damage Restoration (LanPaint K Sampler, 5 steps of thinking)
![Inpainting Result 4](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_07.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_4)
[Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl)
### Example SDXL 5: Huge Damage Restoration (LanPaint K Sampler, 20 steps of thinking)
![Inpainting Result 5](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_08.jpg)
[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_5)
[Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl)
Check more for use cases like inpaint on [fine tuned models](https://github.com/scraed/LanPaint/issues/12#issuecomment-2938662021) and [face swapping](https://github.com/scraed/LanPaint/issues/12#issuecomment-2938723501), thanks to [Amazon90](https://github.com/Amazon90).
## **How to Use These Examples:**
1. Navigate to the **example** folder (i.e example_1) by clicking **View Workflow & Masks**, download all pictures.
2. Drag **InPainted_Drag_Me_to_ComfyUI.png** into ComfyUI to load the workflow.
3. Download the required model from Civitai by clicking **Model Used in This Example**.
4. Load the model into the **"Load Checkpoint"** node.
5. Upload **Original_No_Mask.png** to the **"Load image"** node in the **"Original Image"** group (far left).
6. Upload **Masked_Load_Me_in_Loader.png** to the **"Load image"** node in the **"Mask image for inpainting"** group (second from left).
7. Queue the task, you will get inpainted results from three methods:
- **[VAE Encode for Inpainting](https://comfyanonymous.github.io/ComfyUI_examples/inpaint/)** (middle),
- **[Set Latent Noise Mask](https://comfyui-wiki.com/en/tutorial/basic/how-to-inpaint-an-image-in-comfyui)** (second from right),
- **LanPaint** (far right).
8. You also get an output from "masked blend" node, which copy the original image and paste onto the unmasked part of output. It is useful if you want unmasked region to match original picture pixel perfectly.
Compare and explore the results from each method!
![WorkFlow](https://github.com/scraed/LanPaint/blob/master/Example.JPG)
## Usage
**Workflow Setup**
@@ -71,53 +165,82 @@ Same as default ComfyUI KSampler - simply replace with LanPaint KSampler nodes.
- LanPaint requires binary masks (values of 0 or 1) without opacity or smoothing. To ensure compatibility, set the mask's **opacity and hardness to maximum** in your mask editor. During inpainting, any mask with smoothing or gradients will automatically be converted to a binary mask.
- LanPaint relies heavily on your text prompts to guide inpainting - explicitly describe the content you want generated in the masked area. If results show artifacts or mismatched elements, counteract them with targeted negative prompts.
### Basic Sampler
## Basic Sampler
![Samplers](https://github.com/scraed/LanPaint/blob/master/Nodes.JPG)
**LanPaint KSampler**
- LanPaint KSampler: The most basic and easy to use sampler for inpainting.
- LanPaint KSampler (Advanced): Full control of all parameters.
### LanPaint KSampler
Simplified interface with recommended defaults:
- Steps: 50+ recommended
- LanPaint NumSteps: 1-10 (complexity of edits)
- Built-in parameter presets
- Steps: 20 - 50. More steps will give more "thinking" and better results.
- LanPaint NumSteps: The turns of thinking before denoising. Recommend 5 for most of tasks ( which means 5 times slower than sampling without thinking). Use 10 for more challenging tasks.
- LanPaint Prompt mode: Image First mode and Prompt First mode. Image First mode focuses on the image, inpaint based on image context (maybe ignore prompt), while Prompt First mode focuses more on the prompt. Use Prompt First mode for tasks like character consistency. (Technically, it Prompt First mode change CFG scale to negative value in the BIG score to emphasis prompt, which will costs image quality.)
**LanPaint KSampler (Advanced)**
### LanPaint KSampler (Advanced)
Full parameter control:
## Key Parameters
**Key Parameters**
| Parameter | Range | Description |
|-----------|-------|-------------|
| `Steps` | 0-100 | Total steps of diffusion sampling. Higher means better inpainting. Recommend 50. |
| `LanPaint_NumSteps` | 0-20 | Reasoning iterations per denoising step ("thinking depth"). Easy task: 1-2. Hard task: 5-10 |
| `LanPaint_Lambda` | 0.1-50 | Content alignment strength (higher = stricter). Recommend 6.0 |
| `LanPaint_StepSize` | 0.1-1.0 | The StepSize of each thinking step. Recommend 0.1 for most cases, 0.5 for some difficult cases. |
| `LanPaint_cfg_BIG` | 0-20 | CFG scale used when aligning masked and unmasked region (higher = better alignment). Recommend 8 for seamless inpaint (i.e limbs, faces), 0-1 for character consistency (i.e multiple view) |
| `Steps` | 0-100 | Total steps of diffusion sampling. Higher means better inpainting. Recommend 20-50. |
| `LanPaint_NumSteps` | 0-20 | Reasoning iterations per denoising step ("thinking depth"). Easy task: 2-5. Hard task: 5-10 |
| `LanPaint_Lambda` | 0.1-50 | Content alignment strength (higher = stricter). Recommend 4.0 - 10.0 |
| `LanPaint_StepSize` | 0.1-1.0 | The StepSize of each thinking step. Recommend 0.1-0.5. |
| `LanPaint_Beta` | 0.1-2.0 | The StepSize ratio between masked / unmasked region. Small value can compensate high lambda values. Recommend 1.0 |
| `LanPaint_Friction` | 0.0-100.0 | The friction of Langevin dynamics. Higher means more slow but stable, lower means fast but unstable. Recommend 10.0 - 20.0|
| `LanPaint_EarlyStop` | 0-10 | Stop LanPaint iteration before the final sampling step. Helps to remove artifacts in some cases. Recommend 1-5|
| `LanPaint_PromptMode` | Image First / Prompt First | Image First mode focuses on the image context, maybe ignore prompt. Prompt First mode focuses more on the prompt. |
For detailed descriptions of each parameter, simply hover your mouse over the corresponding input field to view tooltips with additional information.
### LanPaint Mask Blend
This node blends the original image with the inpainted image based on the mask. It is useful if you want the unmasked region to match the original image pixel perfectly.
## LanPaint KSampler (Advanced) Tuning Guide
For challenging inpainting tasks:
1️⃣ **Primary Adjustments**:
- Increase **steps**, **LanPaint_StepSize**,**LanPaint_NumSteps** (thinking iterations), and **LanPaint_cfg_BIG** (guidance scale).
1️⃣ **Boost Quality**
Increase **total number of sampling steps** (very important!), **LanPaint_NumSteps** (thinking iterations) or **LanPaint_Lambda** if the inpainted result does not meet your expectations.
2️⃣ **Secondary Tweaks**:
- Boost **LanPaint_Lambda** (spatial constraint strength) or **LanPaint_StepSize** (denoising aggressiveness).
2️⃣ **Boost Speed**
If you want better results but still need fewer steps, consider:
- **Increasing LanPaint_StepSize** to speed up the thinking process.
- **Decreasing LanPaint_Friction** to make the Langevin dynamics converges more faster.
3️⃣ **Balance Speed vs Stability**:
- Reduce **LanPaint_Friction** to prioritize faster results with fewer "thinking" steps (*may risk instability*).
- Increase **LanPaint_Tamed** (noise normalization onto a sphere) or **LanPaint_Alpha** (constraint the friction of underdamped Langevin dynamics) to suppress artifacts.
3️⃣ **Fix Unstability**:
If you find the results have wired texture, try
- Reduce **LanPaint_Friction** to make the Langevin dynamics more stable.
- Reduce **LanPaint_StepSize** to use smaller step size.
- Reduce **LanPaint_Beta** if you are using a high lambda value.
⚠️ **Notes**:
- Optimal parameters vary depending on the **model** and the **size of the inpainting area**.
- For effective tuning, **fix the seed** and adjust parameters incrementally while observing the results. This helps isolate the impact of each setting.
- For effective tuning, **fix the seed** and adjust parameters incrementally while observing the results. This helps isolate the impact of each setting. Better to do it with a batche of images to avoid overfitting on a single image.
## ToDo
- Try Implement Detailer
- Provide inference code on without GUI.
## Contribute
- 2025/03/06: Bug Fix for str not callable error and unpack error. Big thanks to [jamesWalker55](https://github.com/jamesWalker55) and [EricBCoding](https://github.com/EricBCoding).
Help us improve LanPaint! 🚀 **Report bugs**, share **example cases**, or contribute your **personal parameter settings** to benefit the community.
## 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},
}
```
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.5 MiB

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.3 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.5 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 2.9 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.5 MiB

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.2 MiB

After

Width:  |  Height:  |  Size: 1.2 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.0 MiB

After

Width:  |  Height:  |  Size: 1012 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.2 MiB

After

Width:  |  Height:  |  Size: 1.0 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 856 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 805 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.2 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.3 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.5 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.3 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.5 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.3 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 455 KiB

After

Width:  |  Height:  |  Size: 702 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 428 KiB

After

Width:  |  Height:  |  Size: 644 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 242 KiB

After

Width:  |  Height:  |  Size: 357 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 239 KiB

After

Width:  |  Height:  |  Size: 349 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 256 KiB

After

Width:  |  Height:  |  Size: 379 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 397 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 549 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 606 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 467 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 407 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 782 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 260 KiB

+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "LanPaint"
version = "0.0.8"
version = "1.2.0"
description = "Achieve seamless inpainting results without needing a specialized inpainting model."
authors = [
{name = "LanPaint", email = "czhengac@connect.ust.hk"}
+174
View File
@@ -0,0 +1,174 @@
import torch
from .utils import *
from functools import partial
class LanPaint():
def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False):
self.n_steps = NSteps
self.chara_lamb = Lambda
self.IS_FLUX = IS_FLUX
self.IS_FLOW = IS_FLOW
self.step_size = StepSize
self.inner_model = Model
self.friction = Friction
self.chara_beta = Beta
self.img_dim_size = None
def add_none_dims(self, array):
# Create a tuple with ':' for the first dimension and 'None' repeated num_nones times
index = (slice(None),) + (None,) * (self.img_dim_size-1)
return array[index]
def remove_none_dims(self, array):
# 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 __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 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):
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
if IS_FLUX or IS_FLOW:
x_t = x * ( self.add_none_dims(abt)**0.5 + (1-self.add_none_dims(abt))**0.5 )
else:
x_t = x / ( 1+self.add_none_dims(VE_Sigma)**2 )**0.5 # switch to variance perserving x_t values
############ LanPaint Iterations Start ###############
# after noise_scaling, noise = latent_image + noise * sigma, which is x_t in the variance exploding diffusion model notation for the known region.
args = None
for i in range(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 )
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 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 = out * (1-latent_mask) + self.latent_image * latent_mask
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)
else:
x = x_t * ( 1+sigma**2 )**0.5 # switch to variance exploding
x_0, x_0_BIG = self.inner_model(x, self.remove_none_dims(sigma), model_options=model_options, seed=seed)
score_x = -(x_t - x_0)
score_y = - (1 + lamb) * ( x_t - y ) + lamb * (x_t - x_0_BIG)
return score_x * (1 - mask) + score_y * mask
def sigma_x(self, abt):
# the time scale for the x_t update
return abt**0
def sigma_y(self, abt):
beta = self.chara_beta * abt ** 0
return beta
def langevin_dynamics(self, x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):
# 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)
sigma, abt, dtx, dty, Gamma_x, Gamma_y, A_x, A_y, D_x, D_y = step_sizes
# print('mask',mask.device)
if torch.mean(dtx) <= 0.:
return x_t, args
# -------------------------------------------------------------------------
# Compute the Langevin dynamics update in variance perserving notation
# -------------------------------------------------------------------------
x0 = self.x0_evalutation(x_t, score, sigma, args)
C = abt**0.5 * x0 / (1-abt)
A = A_x * (1-mask) + A_y * mask
D = D_x * (1-mask) + D_y * mask
dt = dtx * (1-mask) + dty * mask
Gamma = Gamma_x * (1-mask) + Gamma_y * mask
def Coef_C(x_t):
x0 = self.x0_evalutation(x_t, score, sigma, args)
C = (abt**0.5 * x0 - x_t )/ (1-abt) + A * x_t
return C
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):
osc = StochasticHarmonicOscillator(Gamma, A, C, D )
x_t, v = osc.dynamics(x_t, v, dt )
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)
C_new = 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
return x_t, (v, C)
def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y):
# -------------------------------------------------------------------------
# Unpack current times parameters (sigma and abt)
sigma, abt, flow_t = current_times
sigma = self.add_none_dims(sigma)
abt = self.add_none_dims(abt)
# 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.
Gamma_hat_x = self.friction **2 * self.step_size * sigma_x / 0.1 * sigma**0
Gamma_hat_y = self.friction **2 * self.step_size * sigma_y / 0.1 * sigma**0
#print("Gamma_hat_x", torch.mean(Gamma_hat_x).item(), "Gamma_hat_y", torch.mean(Gamma_hat_y).item())
# adjust dt to match denoise-addnoise steps sizes
Gamma_hat_x /= 2.
Gamma_hat_y /= 2.
A_t_x = (1) / ( 1 - abt ) * dtx / 2
A_t_y = (1+self.chara_lamb) / ( 1 - abt ) * dty / 2
A_x = A_t_x / (dtx/2)
A_y = A_t_y / (dty/2)
Gamma_x = Gamma_hat_x / (dtx/2)
Gamma_y = Gamma_hat_y / (dty/2)
#D_x = (2 * (1 + sigma**2) )**0.5
#D_y = (2 * (1 + sigma**2) )**0.5
D_x = (2 * 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
+352 -256
View File
@@ -8,15 +8,16 @@ import latent_preview
from functools import partial
from comfy.utils import repeat_to_batch_size
from comfy.samplers import *
# Monkey patch comfy.samplers module by importing with absolute package path
#exec(inspect.getsource(comfy.samplers).replace("from .", "from comfy."))
from comfy.model_base import ModelType
from .utils import *
from .lanpaint import LanPaint
def reshape_mask(input_mask, output_shape):
dims = len(output_shape) - 2
scale_mode = "nearest-exact"
mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode)
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])
@@ -76,10 +77,29 @@ class KSamplerX0Inpaint:
self.sigmas = sigmas
self.model_sigmas = torch.cat( (torch.tensor([0.], device = sigmas.device) , torch.tensor( self.inner_model.model_patcher.get_model_object("model_sampling").sigmas, device = sigmas.device) ) )
self.model_sigmas = torch.tensor( self.model_sigmas, dtype = self.sigmas.dtype )
def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None):
def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None,**kwargs):
### For 1.5 and XL model
# x is x_t in the notation of variance exploding diffusion model, x_t = x_0 + sigma * noise
# sigma is the noise level
# print what is inside model_options
### For flux model
# x is rectified flow x_t = sigma * noise + (1.0 - sigma) * x_0
IS_FLUX = self.inner_model.inner_model.model_type == ModelType.FLUX
IS_FLOW = self.inner_model.inner_model.model_type == ModelType.FLOW
# 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
abt = 1/( 1+VE_Sigma**2 )
Flow_t = (1-abt)**0.5 / ( (1-abt)**0.5 + abt**0.5 )
if denoise_mask is not None:
if "denoise_mask_function" in model_options:
denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
@@ -87,205 +107,31 @@ class KSamplerX0Inpaint:
denoise_mask = (denoise_mask > 0.5).float()
latent_mask = 1 - denoise_mask
current_times = (VE_Sigma, abt, Flow_t)
abt = 1/( 1+sigma**2 )
current_step = torch.argmin( torch.abs( self.sigmas - torch.mean(sigma) ) )
total_steps = len(self.sigmas)-1
if self.step_time_schedule == "dual_shrink":
step_size = self.step_size * (1 - abt) ** 0.5 * abt ** 0.5
elif self.step_time_schedule == "follow_sampler":
time_ind = torch.argmin(torch.abs(self.sigmas - sigma))
times = torch.log( 1+ self.sigmas**2)
time_intervals = times[1:] - times[:-1]
time_intervals = time_intervals / time_intervals[0]
step_size = time_intervals[time_ind] * self.step_size
if total_steps - current_step <= self.LanPaint_early_stop:
out = self.PaintMethod(x, self.latent_image, self.noise, sigma, latent_mask, current_times, model_options, seed, n_steps=0)
else:
step_size = self.step_size * (1 - abt) ** 0.5
current_times = (sigma, abt)
# self.inner_model.inner_model.scale_latent_inpaint returns variance exploding x_t values
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
x_t = x #/ ( 1+sigma**2 )**0.5 # switch to variance perserving x_t values
# after noise_scaling, noise = latent_image + noise * sigma, which is x_t in the variance exploding diffusion model notation for the known region.
args = None
for i in range(self.n_steps):
if sigma > self.start_sigma or sigma < self.end_sigma:
break
score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = abt, sigma = sigma, model_options = model_options, seed = seed )
if self.step_size_schedule == "linear":
step_size_i = step_size * (1 - i/(self.n_steps) )
else:
step_size_i = step_size
x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size_i , current_times, sigma_x = self.sigma_x(abt), sigma_y = self.sigma_y(abt), args = args)
x = x_t #* ( 1+sigma**2 )**0.5
# out is x_0
out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed)
out = out * denoise_mask + self.latent_image * latent_mask
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)
# Only show preview every few steps to improve performance
if current_step % 2 == 0:
# Directly call the preview callback if it exists
callback = model_options.get("callback", None)
if callback is not None:
callback({"i": current_step, "denoised": out, "x": x})
return out
def mid_times(self, current_times, step_size):
sigma, abt = current_times
tt = torch.log(1+sigma**2)
tt_mid = torch.max( tt - step_size, tt*0 )
sigma_mid = (torch.exp(tt_mid) - 1) ** 0.5
sigma_mid_prev = sigma_mid
# find the closest sigma to sigma_mid from self.sigmas
#sigma_mid = self.model_sigmas[torch.argmin(torch.abs(self.model_sigmas - sigma_mid))]
abt_mid = 1/(1+sigma_mid**2)
return sigma_mid, abt_mid
def score_model(self, x_t, y, mask, abt, sigma, model_options, seed):
# the score function for the Langevin dynamics
lamb = self.chara_lamb
beta = self.chara_beta * (1-abt)**0.5
x_0, x_0_BIG = self.inner_model(x_t, sigma, model_options=model_options, seed=seed)
e_t = x_t / ((1 - abt) ** 0.5 * (1 + sigma**2) ** 0.5 )- (abt ** 0.5 / (1 - abt) ** 0.5) * x_0
e_t_BIG = x_t / ((1 - abt) ** 0.5 * (1 + sigma**2) ** 0.5 )- (abt ** 0.5 / (1 - abt) ** 0.5) * x_0_BIG
score_x = -e_t
score_y = - (1 + lamb) * ( x_t/ ((1 + sigma**2) ** 0.5 *(1 - abt)**0.5) - abt**0.5 /(1 - abt)**0.5 * y ) + lamb * e_t_BIG
return score_x * (1 - mask) + score_y * mask
def sigma_x(self, abt):
# the time scale for the x_t update
return 1
def sigma_y(self, abt):
# the time scale for the y_t update
if self.beta_scale == "shrink":
beta = self.chara_beta * (1-abt)**0.5
elif self.beta_scale == "dual_shrink":
beta = self.chara_beta * (1-abt)**0.5 * abt ** 0.5
elif self.beta_scale == "back_shrink":
beta = self.chara_beta * abt ** 0.5
else:
beta = self.chara_beta
return beta
def langevin_dynamics(self, x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):
# -------------------------------------------------------------------------
# Unpack current times parameters (sigma and abt)
sigma, abt = current_times
# Compute time step (dtx, dty) for x and y branches.
dtx = 2 * step_size * sigma_x
dty = 2 * step_size * sigma_y
if self.step_time_schedule == "dual_shrink":
ref_dt = 0.1 * (1-abt)**0.5 * abt ** 0.5
else:
ref_dt = 0.1 * (1-abt)**0.5
# -------------------------------------------------------------------------
# Define friction parameter Gamma_hat for each branch.
# Using dtx**0 provides a tensor of the proper device/dtype.
Gamma_hat_x = self.friction * dtx / (1e-4+ 2 * sigma_x * ref_dt)
Gamma_hat_y = self.friction * dty / (1e-4+ 2 * sigma_y * ref_dt)
# Get mid time parameters (sigma_mid and abt_mid) for each branch.
sigma_mid_x, abt_mid_x = self.mid_times(current_times, dtx)
sigma_mid_y, abt_mid_y = self.mid_times(current_times, dty)
if sigma_mid_x >= sigma or sigma_mid_y >= sigma:
return x_t, args
# -------------------------------------------------------------------------
# A: Update epsilon (score estimate and noise initialization)
# -------------------------------------------------------------------------
# Compute the score-based epsilon (scaled as sqrt(1-abt))
score_model = score(x_t)
eps_model = -score_model
# Initialize epsilon and Z if not provided in args.
if args is None:
eps = eps_model
Z = torch.randn_like(x_t)
else:
eps, Z = args
# -------------------------------------------------------------------------
# B: Update epsilon mean dynamics and compute the mid-point in z-space.
# -------------------------------------------------------------------------
# Compute the weighted combination term for epsilon mean update:
# term = (2/Γ_hat)*(1-exp(-0.5*Γ_hat))
term_x = 2.0 / (Gamma_hat_x + 1e-4) * (1 - torch.exp(-0.5 * Gamma_hat_x))
term_y = 2.0 / (Gamma_hat_y + 1e-4) * (1 - torch.exp(-0.5 * Gamma_hat_y))
eps_bar_x = term_x * eps + (1 - term_x) * eps_model
eps_bar_y = term_y * eps + (1 - term_y) * eps_model
# Combine branches according to mask.
eps_bar = eps_bar_x * (1 - mask) + eps_bar_y * mask
# Form the denoised epsilon using self.alpha (assumed to be 1/Ψ)
eps_denoise = self.alpha * eps_bar + (1 - self.alpha) * eps_model
# tamed
eps_model_x = eps_denoise* (1 - mask)
eps_model_x = eps_model_x* (torch.sum(1 - mask, dim = (1,2,3))/torch.sum(eps_model_x**2, dim = (1,2,3))) **0.5 ** torch.minimum(self.tamed*(dtx),sigma**0)#/( 1 + self.tamed*(sigma - sigma_mid_x) * (torch.sum(eps_model_x**2)/torch.sum((1 - mask)))**0.5 )
eps_model_y = eps_denoise* mask
eps_model_y = eps_model_y* (torch.sum(mask, dim = (1,2,3))/torch.sum(eps_model_y**2, dim = (1,2,3))) **0.5 ** torch.minimum(self.tamed*(dty),sigma**0)#/( 1 + self.tamed*(sigma - sigma_mid_y) * (torch.sum(eps_model_y**2)/torch.sum(mask))**0.5 )
eps_denoise = eps_model_x * (1 - mask) + eps_model_y * mask
# Update the mean epsilon for the next step:
eps_x = eps * torch.exp(-0.5 * Gamma_hat_x) + eps_model * (1 - torch.exp(-0.5 * Gamma_hat_x))
eps_y = eps * torch.exp(-0.5 * Gamma_hat_y) + eps_model * (1 - torch.exp(-0.5 * Gamma_hat_y))
eps = eps_x * (1 - mask) + eps_y * mask
# Transform x to z using z = x * sqrt(1+sigma^2). Here we have already set x to z to avoid floating point stability issue.
z_t = x_t #* (1 + sigma**2) ** 0.5
# Compute the mid-point update in z-space for each branch:
z_mid_x = z_t + eps_denoise * (sigma_mid_x - sigma)
z_mid_y = z_t + eps_denoise * (sigma_mid_y - sigma)
z_mid = z_mid_x * (1 - mask) + z_mid_y * mask
# -------------------------------------------------------------------------
# C: Update noise terms and finalize the x update.
# -------------------------------------------------------------------------
# Generate auxiliary noise terms.
Z_q = torch.randn_like(x_t)
Z_q_avg = torch.randn_like(x_t)
Z_z = torch.randn_like(x_t)
# Update Z for each branch:
Z_x = torch.exp(-0.5 * Gamma_hat_x) * Z + (1 - torch.exp(-Gamma_hat_x)) ** 0.5 * Z_q
Z_y = torch.exp(-0.5 * Gamma_hat_y) * Z + (1 - torch.exp(-Gamma_hat_y)) ** 0.5 * Z_q
Z_next = Z_x * (1 - mask) + Z_y * mask
# Compute the combined noise update following the scheme:
Z_comb_x = (
(1 - torch.exp(-Gamma_hat_x / 2)) / torch.sqrt(Gamma_hat_x + 1e-4) *
(Z + torch.sqrt(torch.tanh(Gamma_hat_x / 4)) * Z_q)
+ torch.sqrt(1 - (4 / (Gamma_hat_x + 1e-4)) * torch.tanh(Gamma_hat_x / 4)) * Z_q_avg
)
Z_comb_y = (
(1 - torch.exp(-Gamma_hat_y / 2)) / torch.sqrt(Gamma_hat_y + 1e-4) *
(Z + torch.sqrt(torch.tanh(Gamma_hat_y / 4)) * Z_q)
+ torch.sqrt(1 - (4 / (Gamma_hat_y + 1e-4)) * torch.tanh(Gamma_hat_y / 4)) * Z_q_avg
)
Z_comb = Z_comb_x * (1 - mask) + Z_comb_y * mask
# Combine with an additional noise term using self.alpha.
Z_comb = self.alpha ** 0.5 * Z_comb + (1 - self.alpha) ** 0.5 * Z_z
# Compute the change in sigma (dsigma = sqrt(sigma^2 - sigma_mid^2)).
dsigma_x = sigma * torch.sqrt(1 - (sigma_mid_x / sigma) ** 2)
dsigma_y = sigma * torch.sqrt(1 - (sigma_mid_y / sigma) ** 2)
dsigma = dsigma_x * (1 - mask) + dsigma_y * mask
# Final z update.
z_final = z_mid + Z_comb * dsigma
# Transform back to x-space: x = z / sqrt(1+sigma^2)
x_t = z_final #/ (1 + sigma**2) ** 0.5
return x_t, (eps, Z_next)
# Custom sampler class extending ComfyUI's KSAMPLER for LanPaint
class KSAMPLER(comfy.samplers.KSAMPLER):
def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
@@ -299,30 +145,38 @@ class KSAMPLER(comfy.samplers.KSAMPLER):
model_k.noise = torch.randn(noise.shape, generator=generator, device="cpu").to(noise.dtype).to(noise.device)
else:
model_k.noise = noise
model_wrap.cfg_BIG = model_wrap.model_patcher.LanPaint_cfg_BIG
model_k.step_size = model_wrap.model_patcher.LanPaint_StepSize
model_k.chara_lamb = model_wrap.model_patcher.LanPaint_Lambda
model_k.chara_beta = model_wrap.model_patcher.LanPaint_Beta
model_k.n_steps = model_wrap.model_patcher.LanPaint_NumSteps
model_k.friction = model_wrap.model_patcher.LanPaint_Friction
model_k.alpha = model_wrap.model_patcher.LanPaint_Alpha
model_k.tamed = model_wrap.model_patcher.LanPaint_Tamed
model_k.beta_scale = model_wrap.model_patcher.LanPaint_BetaScale
model_k.step_size_schedule = model_wrap.model_patcher.LanPaint_StepSizeSchedule
model_k.step_time_schedule = model_wrap.model_patcher.LanPaint_StepTimeSchedule
model_k.start_sigma = model_wrap.model_patcher.LanPaint_StartSigma
model_k.end_sigma = model_wrap.model_patcher.LanPaint_EndSigma
IS_FLUX = model_wrap.inner_model.model_type == ModelType.FLUX
IS_FLOW = model_wrap.inner_model.model_type == ModelType.FLOW
# unify the notations into variance exploding diffusion model
if IS_FLUX:
model_wrap.cfg_BIG = 1.0
else:
model_wrap.cfg_BIG = model_wrap.model_patcher.LanPaint_cfg_BIG
noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
model_k.PaintMethod = LanPaint(model_k.inner_model,
model_wrap.model_patcher.LanPaint_NumSteps,
model_wrap.model_patcher.LanPaint_Friction,
model_wrap.model_patcher.LanPaint_Lambda,
model_wrap.model_patcher.LanPaint_Beta,
model_wrap.model_patcher.LanPaint_StepSize,
IS_FLUX = IS_FLUX,
IS_FLOW = IS_FLOW)
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.
k_callback = None
total_steps = len(sigmas) - 1
if callback is not None:
k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
print("LanPaint KSampler call sampler_function", self.sampler_function)
#print("LanPaint KSampler call sampler_function", self.sampler_function)
# The main loop!
#print("##########")
#print("Sampling with ", self.sampler_function)
#print("##########")
samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
print("LanPaint KSampler end sampler_function")
#print("LanPaint KSampler end sampler_function")
samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
return samples
@@ -345,9 +199,33 @@ def override_sample_function():
comfy.samplers.CFGGuider.outer_sample = original_outer_sample
class LanPaint_UpSale_LatentNoiseMask:
@classmethod
def INPUT_TYPES(s):
return {"required": { "samples": ("LATENT",),
"scale": ("INT", {"default": 2, "min": 2, "max": 8, "step": 1}),
}}
RETURN_TYPES = ("LATENT",)
FUNCTION = "set_mask"
KSAMPLER_NAMES = ["euler", "dpmpp_2m", "uni_pc"]
CATEGORY = "latent/inpaint"
def set_mask(self, samples, scale):
s = samples.copy()
samples = s['samples']
# generate a mask with every scaleth pixel set to 1
mask = torch.zeros(samples.shape[0], 1, samples.shape[2], samples.shape[3], device=samples.device) + 1
mask[:, :, ::scale, ::scale] = 0
s["noise_mask"] = mask
return (s,)
#KSAMPLER_NAMES = ["euler", "dpmpp_2m", "uni_pc"]
KSAMPLER_NAMES = ["euler","euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
"dpm_fast", "dpmpp_sde", "dpmpp_sde_gpu",
"dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm",
"deis", "res_multistep", "res_multistep_ancestral",
"gradient_estimation", "er_sde", "seeds_2", "seeds_3"]
class LanPaint_KSampler():
@classmethod
@@ -356,16 +234,17 @@ class LanPaint_KSampler():
"required": {
"model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}),
"steps": ("INT", {"default": 50, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
"steps": ("INT", {"default": 30, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
"cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}),
"sampler_name": (KSAMPLER_NAMES, {"tooltip": "Recommended: euler."}),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "karras", "tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
"positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}),
"negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
"latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling."}),
"LanPaint_NumSteps": ("INT", {"default": 10, "min": 0, "max": 20, "tooltip": "The number of steps for the Langevin dynamics, representing the turns of thinking per step."}),
"LanPaint_Info": ("STRING", {"default": "LanPaint KSampler. Recommend steps 50 ( increase steps boosts performance ), LanPaint NumSteps 1-10 depending on the difficulty of task. For more information, visit https://github.com/scraed/LanPaint", "multiline": True}),
"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}),
}
}
@@ -376,20 +255,18 @@ 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_StepSize=0.05, LanPaint_NumSteps=5, LanPaint_Info=""):
model.LanPaint_StepSize = 0.3
model.LanPaint_Lambda = 6.0
model.LanPaint_Beta = 0.6
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=""):
model.LanPaint_StepSize = 0.15
model.LanPaint_Lambda = 16.0
model.LanPaint_Beta = 1.
model.LanPaint_NumSteps = LanPaint_NumSteps
model.LanPaint_Friction = 10.
model.LanPaint_Alpha = 0.5
model.LanPaint_Tamed = 0.1
model.LanPaint_BetaScale = "shrink"
model.LanPaint_StepSizeSchedule = "linear"
model.LanPaint_StepTimeSchedule = "shrink"
model.LanPaint_StartSigma = 20.
model.LanPaint_EndSigma = 1.
model.LanPaint_cfg_BIG = cfg
model.LanPaint_Friction = 15.
model.LanPaint_EarlyStop = 1
if LanPaint_PromptMode == "Image First":
model.LanPaint_cfg_BIG = cfg
else:
model.LanPaint_cfg_BIG = 0*cfg - 0.5
with override_sample_function():
return nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
class LanPaint_KSamplerAdvanced:
@@ -399,7 +276,7 @@ class LanPaint_KSamplerAdvanced:
{"model": ("MODEL",),
"add_noise": (["enable", "disable"], ),
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 50, "min": 1, "max": 10000}),
"steps": ("INT", {"default": 30, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
"sampler_name": (KSAMPLER_NAMES, ),
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
@@ -409,20 +286,14 @@ class LanPaint_KSamplerAdvanced:
"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"], ),
"LanPaint_NumSteps": ("INT", {"default": 10, "min": 0, "max": 20, "tooltip": "The number of steps for the Langevin dynamics, representing the turns of thinking per step."}),
"LanPaint_Lambda": ("FLOAT", {"default": 6., "min": 0.1, "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The lambda parameter for the bidirectional guidance. Higher values align with known regions more closely, but may result in instability."}),
"LanPaint_StepSize": ("FLOAT", {"default": 0.3, "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": 0.6, "min": 0.0001, "max": 5, "step": 0.1, "round": 0.1, "tooltip": "The beta parameter for the bidirectional guidance. Scale the step size for the known region independently for the Langevin dynamics. Higher values result in faster convergence but may be unstable."}),
"LanPaint_Friction": ("FLOAT", {"default": 10., "min": 1., "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The friction parameter for the underdamped Langevin dynamics, higher values result in faster convergence but may be unstable."}),
"LanPaint_Alpha": ("FLOAT", {"default": 0.5, "min": 0.0001, "max": 1., "step": 0.1, "round": 0.1, "tooltip": "The (rescaled) alpha parameter for the HFHR langevin dynamics, mixes Langevin dynamics and underdamped Langevin dynamics with a friction term. 0 corresponds to Langevin dynamics, 1 corresponds to underdamped Langevin dynamics."}),
"LanPaint_Tamed": ("FLOAT", {"default": 1., "min": 0.000, "max": 20., "step": 0.1, "round": 0.1, "tooltip": "The tame strength for the noise, normalize and projects the noise onto unit sphere to enhance stability."}),
"LanPaint_BetaScale": (["shrink", "fixed", "dual_shrink", "back_shrink"], {"default": "shrink", "tooltip": "The beta scale, determines how the beta parameter changes over time. Shrink: beta = beta * (1 - alpha bar) ** 0.5; Fixed: beta = beta; Dual_shrink: beta = beta * (1 - alpha bar) ** 0.5 * alpha bar ** 0.5; Back_shrink: beta = beta * alpha bar ** 0.5; Alpha bar: the alpha cumprod."}),
"LanPaint_StepSizeSchedule": (["const", "linear"], {"default": "linear", "tooltip": "The step size schedule for the Langevin dynamics, const: constant step size, linear: linearly decreasing step size."}),
"LanPaint_StepTimeSchedule": (["shrink", "dual_shrink", "follow_sampler"], {"default": "shrink", "tooltip": "The step size schedule for the first step of Langevin dynamics during diffusion sampling, shrink: step size = step size * (1 - alpha bar) ** 0.5; Dual_shrink: step size = step size * (1 - alpha bar) ** 0.5 * alpha bar ** 0.5; Follow_sampler: scale with the sampler step size."}),
"LanPaint_StartSigma": ("FLOAT", {"default": 20., "min": 0.0001, "max": 20.0, "step": 0.1, "round": 0.1, "tooltip": "Start 'thinking' with Langevin dynamics at this sigma value."}),
"LanPaint_EndSigma": ("FLOAT", {"default": 1., "min": 0.000, "max": 20.0, "step": 0.1, "round": 0.1, "tooltip": "Stop 'thinking' with Langevin dynamics at this sigma value."}),
"LanPaint_cfg_BIG": ("FLOAT", {"default": 8., "min": 0., "max": 20.0, "step": 0.1, "round": 0.1, "tooltip": "The CFG scale used in the bidirectional guidance (for the known region only). Higher value results in more closely matching the known region."}),
"LanPaint_Info": ("STRING", {"default": "LanPaint KSampler Advanced. For difficult tasks, first try increasing steps, LanPaint_NumSteps, and LanPaint_cfg_BIG. Then try increase LanPaint_Lambda or LanPaint_StepSize. Decrease LanPaint_Friction if you want to obtain good results with fewer turns of thinking (LanPaint_NumSteps) at the risk of irregular behavior. Increase LanPaint_Tamed or LanPaint_Alpha can suppress irregular behavior. For more information, visit https://github.com/scraed/LanPaint", "multiline": True}),
"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_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}),
},
}
@@ -431,7 +302,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_Alpha=1, LanPaint_Tamed=0., LanPaint_BetaScale="fixed", LanPaint_StepSizeSchedule = "const", LanPaint_StepTimeSchedule = "shrink", LanPaint_StartSigma=20, LanPaint_EndSigma=0, LanPaint_cfg_BIG = 5., 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, 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=""):
force_full_denoise = True
if return_with_leftover_noise == "enable":
force_full_denoise = False
@@ -443,27 +314,252 @@ class LanPaint_KSamplerAdvanced:
model.LanPaint_Beta = LanPaint_Beta
model.LanPaint_NumSteps = LanPaint_NumSteps
model.LanPaint_Friction = LanPaint_Friction
model.LanPaint_Alpha = LanPaint_Alpha
model.LanPaint_Tamed = LanPaint_Tamed
model.LanPaint_BetaScale = LanPaint_BetaScale
model.LanPaint_StepSizeSchedule = LanPaint_StepSizeSchedule
model.LanPaint_StepTimeSchedule = LanPaint_StepTimeSchedule
model.LanPaint_StartSigma = LanPaint_StartSigma
model.LanPaint_EndSigma = LanPaint_EndSigma
model.LanPaint_cfg_BIG = LanPaint_cfg_BIG
model.LanPaint_EarlyStop = LanPaint_EarlyStop
if LanPaint_PromptMode == "Image First":
model.LanPaint_cfg_BIG = cfg
else:
model.LanPaint_cfg_BIG = 0*cfg - 0.5
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)
class MaskBlend:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image1": ("IMAGE", {"tooltip": "Image before inpaint"}),
"image2": ("IMAGE", {"tooltip": "Image after inpaint"}),
"mask": ("MASK",),
"blend_overlap": ("INT", {"default": 1, "min": 1, "max": 51, "step": 2, "tooltip": "The number of pixels to blend between the two images."})
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "blend_images"
CATEGORY = "image/postprocessing"
def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, mask: torch.Tensor, blend_overlap: int):
# smooth the binary 01 mask, keep 1 still 1, but smooth the transition from 1 to 0
# 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.")
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,:,:]
blended_image = image1 * (1 - mask[...,None]) + image2 * mask[...,None]
return (blended_image,)
def gaussian_kernel(self,kernel_size):
"""
Creates a 2D Gaussian kernel with the given size and standard deviation (sigma).
"""
sigma = (kernel_size - 1)/4
# Create a grid of (x, y) coordinates
x = torch.arange(kernel_size).float() - kernel_size // 2
y = torch.arange(kernel_size).float() - kernel_size // 2
x_grid, y_grid = torch.meshgrid(x, y, indexing='ij')
# Compute the Gaussian function
kernel = torch.exp(-(x_grid ** 2 + y_grid ** 2) / (2 * sigma ** 2))
kernel = kernel / kernel.sum() # Normalize the kernel
return kernel
class Noise_EmptyNoise:
def generate_noise(self, latent):
return torch.zeros_like(latent["samples"])
class Noise_RandomNoise:
def __init__(self, seed):
self.seed = seed
def generate_noise(self, latent):
torch.manual_seed(self.seed)
return torch.randn_like(latent["samples"])
# Custom sampler implementation mimmicking base comfy nodes_custom_sampler.py
class LanPaint_SamplerCustom:
@classmethod
def INPUT_TYPES(s):
return {"required":
{"model": ("MODEL",),
"add_noise": ("BOOLEAN", {"default": True}),
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"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_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: prioritizes image quality; Prompt First: prioritizes prompt adherence."}),
"LanPaint_Info": ("STRING", {"default": "LanPaint Custom Sampler. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
}
}
RETURN_TYPES = ("LATENT", "LATENT")
RETURN_NAMES = ("output", "denoised_output")
FUNCTION = "sample"
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_Lambda = 16.0
model.LanPaint_Beta = 1.
model.LanPaint_NumSteps = LanPaint_NumSteps
model.LanPaint_Friction = 15.
model.LanPaint_EarlyStop = 1
if LanPaint_PromptMode == "Image First":
model.LanPaint_cfg_BIG = cfg
else:
model.LanPaint_cfg_BIG = 0 * cfg - 0.5
with override_sample_function():
latent = latent_image.copy()
latent_image = latent["samples"]
latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
latent["samples"] = latent_image
if not add_noise:
noise = Noise_EmptyNoise().generate_noise(latent)
else:
noise = Noise_RandomNoise(noise_seed).generate_noise(latent)
noise_mask = None
if "noise_mask" in latent:
noise_mask = latent["noise_mask"]
x0_output = {}
callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image,noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
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())
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"], ),
"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_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}),
}
}
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
model = guider.model_patcher
model.LanPaint_StepSize = LanPaint_StepSize
model.LanPaint_Lambda = LanPaint_Lambda
model.LanPaint_Beta = LanPaint_Beta
model.LanPaint_NumSteps = LanPaint_NumSteps
model.LanPaint_Friction = LanPaint_Friction
model.LanPaint_EarlyStop = LanPaint_EarlyStop
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
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)
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 = 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())
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 = {
"LanPaint_KSampler": LanPaint_KSampler,
"LanPaint_KSamplerAdvanced": LanPaint_KSamplerAdvanced,
"LanPaint_SamplerCustom" : LanPaint_SamplerCustom,
"LanPaint_SamplerCustomAdvanced" : LanPaint_SamplerCustomAdvanced,
"LanPaint_MaskBlend": MaskBlend,
# "LanPaint_UpSale_LatentNoiseMask": LanPaint_UpSale_LatentNoiseMask,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
"LanPaint_KSampler": "LanPaint KSampler",
"LanPaint_KSamplerAdvanced": "LanPaint KSampler (Advanced)"
"LanPaint_KSamplerAdvanced": "LanPaint KSampler (Advanced)",
"LanPaint_SamplerCustom" : "LanPaint Sampler Custom",
"LanPaint_SamplerCustomAdvanced" : "LanPaint Sampler Custom (Advanced)",
"LanPaint_MaskBlend": "LanPaint Mask Blend",
# "LanPaint_UpSale_LatentNoiseMask": "LanPaint UpSale Latent Noise Mask"
}
+301
View File
@@ -0,0 +1,301 @@
import torch
def epxm1_x(x):
# Compute the (exp(x) - 1) / x term with a small value to avoid division by zero.
result = torch.special.expm1(x) / x
# replace NaN or inf values with 0
result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
mask = torch.abs(x) < 1e-2
result = torch.where(mask, 1 + x/2. + x**2 / 6., result)
return result
def epxm1mx_x2(x):
# Compute the (exp(x) - 1 - x) / x**2 term with a small value to avoid division by zero.
result = (torch.special.expm1(x) - x) / x**2
# replace NaN or inf values with 0
result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
mask = torch.abs(x**2) < 1e-2
result = torch.where(mask, 1/2. + x/6 + x**2 / 24 + x**3 / 120, result)
return result
def expm1mxmhx2_x3(x):
# Compute the (exp(x) - 1 - x - x**2 / 2) / x**3 term with a small value to avoid division by zero.
result = (torch.special.expm1(x) - x - x**2 / 2) / x**3
# replace NaN or inf values with 0
result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
mask = torch.abs(x**3) < 1e-2
result = torch.where(mask, 1/6 + x/24 + x**2 / 120 + x**3 / 720 + x**4 / 5040, result)
return result
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
"""
# Main computation
is_positive = delta > 0
sqrt_abs_delta = torch.sqrt(torch.abs(delta))
gamma_t_sqrt_delta = gamma_t * sqrt_abs_delta
numerator_pos = torch.exp(-gamma_t) - (torch.exp(gamma_t * (sqrt_abs_delta - 1)) + torch.exp(gamma_t * (-sqrt_abs_delta - 1))) / 2
numerator_neg = torch.exp(-gamma_t) * ( 1 - torch.cos(gamma_t * sqrt_abs_delta ) )
numerator = torch.where(is_positive, numerator_pos, numerator_neg)
result = numerator / (delta * gamma_t**2 )
# Handle NaN/inf cases
result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
# Handle numerical instability for small delta
mask = torch.abs(gamma_t_sqrt_delta**2) < 5e-2
taylor = ( -0.5 - gamma_t**2 / 24 * delta - gamma_t**4 / 720 * delta**2 ) * torch.exp(-gamma_t)
result = torch.where(mask, taylor, result)
return result
def exp_sinh_GsqrtD(gamma_t, delta):
"""
Compute e^(-Γt) * sinh(Γt√Δ) / (Γt√Δ)
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
"""
# Main computation
is_positive = delta > 0
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))
# Taylor expansion for small gamma_t_sqrt_delta
mask = torch.abs(gamma_t_sqrt_delta) < 1e-2
taylor = ( 1 + gamma_t**2 / 6 * delta + gamma_t**4 / 120 * delta**2 ) * torch.exp(-gamma_t)
result_pos = torch.where(mask, taylor, result_pos)
# Handle negative delta
result_neg = torch.exp(-gamma_t) * torch.special.sinc(gamma_t_sqrt_delta/torch.pi)
result = torch.where(is_positive, result_pos, result_neg)
return result
def exp_cosh(gamma_t, delta):
"""
Compute e^(-Γt) * cosh(Γt√Δ)
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
"""
exp_1mcosh_GD_result = exp_1mcosh_GD(gamma_t, delta) # e^(-Γt) * (1 - cosh(Γt√Δ))/ ( (Γt)**2 Δ )
result = torch.exp(-gamma_t) - gamma_t**2 * delta * exp_1mcosh_GD_result
return result
def exp_sinh_sqrtD(gamma_t, delta):
"""
Compute e^(-Γt) * sinh(Γt√Δ) / √Δ
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
"""
exp_sinh_GsqrtD_result = exp_sinh_GsqrtD(gamma_t, delta) # e^(-Γt) * sinh(Γt√Δ) / (Γt√Δ)
result = gamma_t * exp_sinh_GsqrtD_result
return result
def zeta1(gamma_t, delta):
# Compute hyperbolic terms and exponential
half_gamma_t = gamma_t / 2
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)
term2 = epxm1mx_x2(-gamma_t)
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
"""
exp_term = torch.exp(-gamma_t)
# 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)
taylor = (
gamma_t*exp_1mcosh_GD_term + 0.5 * gamma_t * exp_sinh_GsqrtD(gamma_t, delta**0)
- 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
def zeta2(gamma_t, delta):
half_gamma_t = gamma_t / 2
return exp_sinh_GsqrtD(half_gamma_t, delta)
def sig11(gamma_t, delta):
return 1 - torch.exp(-gamma_t) + gamma_t**2 * exp_1mcosh_GD(gamma_t, delta) + exp_sinh_sqrtD(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
Zcoef2 = Zcoef1 * gamma_t *( - 2 * exp_1mcosh_GD(gamma_t, delta) / sig11(gamma_t, delta) ) ** 0.5
#cterm = exp_cosh_minus_terms(gamma_t, delta)
#sterm = exp_sinh_sqrtD(gamma_t, delta**0) + exp_sinh_sqrtD(gamma_t, delta)
#Zcoef3 = 2 * torch.sqrt( cterm / ( gamma_t * (1 - delta) * cterm + sterm ) )
Zcoef3 = torch.sqrt( torch.maximum(1 - Zcoef1**2 - Zcoef2**2, sq_total.new_zeros(sq_total.shape)) )
return Zcoef1 * amplitude, Zcoef2 * amplitude, Zcoef3 * amplitude, amplitude
def Zcoefs_asymp(gamma_t, delta):
A_t = (gamma_t * (1 - delta) )/4
return epxm1_x(- 2 * A_t)
class StochasticHarmonicOscillator:
"""
Simulates a stochastic harmonic oscillator governed by the equations:
dy(t) = q(t) dt
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
Γ - Damping coefficient
A - Harmonic potential strength
C - Constant force term
D - Noise amplitude
dw(t) - Wiener process (Brownian motion)
"""
def __init__(self, Gamma, A, C, D):
self.Gamma = Gamma
self.A = A
self.C = C
self.D = D
self.Delta = 1 - 4 * A / Gamma
def sig11(self, gamma_t, delta):
return 1 - torch.exp(-gamma_t) + gamma_t**2 * exp_1mcosh_GD(gamma_t, delta) + exp_sinh_sqrtD(gamma_t, delta)
def sig22(self, gamma_t, delta):
return 1- zeta1(2*gamma_t, delta) + 2 * gamma_t * exp_1mcosh_GD(gamma_t, delta)
def dynamics(self, y0, v0, t):
"""
Calculates the position and velocity variables at time t.
Parameters:
y0 (float): Initial position
v0 (float): Initial velocity v(0) = q(0) / √Γ
t (float): Time at which to evaluate the dynamics
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
A = self.A + dummyzero
C = self.C + dummyzero
D = self.D + dummyzero
Gamma = self.Gamma + dummyzero
zeta_1 = zeta1( Gamma_hat, Delta)
zeta_2 = zeta2( Gamma_hat, Delta)
EE = 1 - Gamma_hat * zeta_2
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)
# sample new position and velocity with multivariate normal distribution
batch_shape = y0.shape
cov_matrix = torch.zeros(*batch_shape, 2, 2, device=y0.device, dtype=y0.dtype)
cov_matrix[..., 0, 0] = cov_yy
cov_matrix[..., 0, 1] = cov_yv
cov_matrix[..., 1, 0] = cov_yv # symmetric
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)
tol = 1e-8
cov_yy = torch.clamp( cov_yy, min = tol )
sd_yy = torch.sqrt( cov_yy )
inv_sd_yy = 1/(sd_yy)
scale_tril[..., 0, 0] = sd_yy
scale_tril[..., 0, 1] = 0.
scale_tril[..., 1, 0] = cov_yv * inv_sd_yy
scale_tril[..., 1, 1] = torch.clamp( cov_vv - cov_yv**2 / cov_yy, min = tol ) ** 0.5
# check if it matches torch.linalg.
#assert torch.allclose(torch.linalg.cholesky(cov_matrix), scale_tril, atol = 1e-4, rtol = 1e-4 )
# Sample correlated noise from multivariate normal
mean = torch.zeros(*batch_shape, 2, device=y0.device, dtype=y0.dtype)
mean[..., 0] = y_mean
mean[..., 1] = v_mean
new_yv = torch.distributions.MultivariateNormal(
loc=mean,
scale_tril=scale_tril
).sample()
return new_yv[...,0], new_yv[...,1]