levzalt/Nanosaur2-Inpaint-ControlNet

Nanosaur2 Inpaint ControlNet

A Nanosaur2-specific, LLLite-style inpainting adapter, with the updated ComfyUI node package needed to run it. The adapter has 8,898,496 parameters and is approximately 35.6 MB. It uses the original Nanosaur2-670M DiT, text encoder, and VAE. Those base weights are not included here.

This is the 1,000-step checkpoint from the 1,000-image continuation run, selected by visual preference. It is an experimental adapter: broad structure can work well, but mask seams, style shifts, changed anatomy, and duplicated objects remain possible. Large masks can substantially change the composition.

Download and install

1. Download nanosaur2_inpaint_step1000.safetensors to ComfyUI/models/controlnet/. The node also supports models/model_patches/.

2. Copy this repository's entire nanosaur2_support folder into ComfyUI/custom_nodes/. Replace the existing Nanosaur2 package with this version; the adapter requires its updated model code as well as the new apply node. Do not install duplicate copies of the package under other names.

3. Download the three base files from well9472/Nanosaur2-670M:

   File	ComfyUI folder
   nanosaur2_diffusion_model.safetensors	models/diffusion_models/
   nanosaur2_text_encoder.safetensors	models/text_encoders/
   nanosaur2_vae.safetensors	models/vae/

4. Restart ComfyUI. Load nanosaur2_inpaint_workflow.json.

5. Select your image in Load Image, paint the repaint region in its mask editor, and edit the positive prompt to describe the intended image. Set Empty Latent Image to the source image's pixel dimensions, using multiples of 16. The template defaults to 896×1152. Resize the source first if necessary, keeping the image and mask aligned.

Use a recent ComfyUI version. The node package uses ComfyUI's model patcher, dynamic-loading, and optimized-operation APIs; older installations may require an update. No additional Python dependencies are introduced by the adapter.

The API workflow is nanosaur2_inpaint_api.json. Both workflows are templates: choose_your_image.png is a placeholder, and no example input image is distributed.

Workflow

Load Nanosaur2 → Apply Nanosaur2 Inpaint Control → KSampler → VAE Decode
                          ↑                                     ↓
                     Image + Mask                 Image Composite Masked
                                                        ↓
                                                    Save Image

Select controlnet/nanosaur2_inpaint_step1000.safetensors in the apply node. White mask pixels mean repaint; black pixels mean preserve.

The supplied workflow composites the decoded result onto the original image using the same mask. That final step preserves the unmasked source pixels. The raw decoded prediction alone does not guarantee pixel preservation.

Setting	Reference value
Sampler / scheduler	Euler / simple
Steps	50
CFG	4
Loader guidance	alternate (CFG / path-drop guidance)
Flow shift	3, supplied by the loader
Control strength	1.0
Control start / end	0.0 / 1.0
Denoise	1.0
Positive prefix	newest, masterpiece
Negative prompt	oldest, low quality

Start from an empty latent, as in the template. Current ComfyUI converts the standard Empty Latent Image output to Nanosaur2's 64 latent channels and 16× spatial compression during sampling. VAE Encode for Inpainting and latent noise-mask blending would change the sampling procedure used in these tests.

Architecture and conditioning

This is inspired by LLLite, but it is not weight-compatible with Anima or SDXL LLLite, and cannot be loaded through their apply nodes.

• Input: four channels, consisting of RGB in [-1,1] and a binary mask in [0,1]. RGB values inside the mask are filled with -1.
• Conditioning encoder: two stride-4 convolutions, four residual blocks, 128 feature channels, and LayerNorm. Total spatial downsampling is 16×.
• Rank-64 residual adapters act before each of the 18 DiT blocks' image QKV projection and MLP input projection.
• Only the adapter was trained; the base DiT, text encoder, and VAE were frozen.

The ComfyUI implementation handles combined positive/negative batches and selects the matching feature rows when path-drop guidance skips middle blocks. Adapter weights use an additional ModelPatcher for device/offload management. Conditioning is scoped to each forward, with no persistent hooks on the shared base model.

Training and validation

The initial experiment used 500 images (450 training / 50 validation), 2,000 steps at 512px area, and learning rate 3e-4. The dataset was expanded to 1,000 images while retaining the same validation split: 950 training / 50 validation. The continuation used 1024px aspect-ratio buckets, learning rate 1e-4, batch size 1, gradient accumulation 4, and a fresh optimizer with a 100-step warmup. This release is continuation step 1,000, following the earlier 2,000-step pilot.

Clean latents and text conditioning were cached. Rectangle, brush, and edge masks were generated dynamically; masked RGB controls were not precached. The training images were a fixed sample from two local illustration collections; the dataset is not included in this release.

The checkpoint was selected visually, rather than by minimum validation loss. ComfyUI integration checks covered numerical agreement with the standalone adapter, mixed CFG/path-drop batches, mask isolation, zero strength, timestep gating, and isolation from unpatched model branches. A full 896×1152, 50-step ComfyUI workflow passed, including byte-identical unmasked pixels after compositing. These are implementation checks, not a broad quality benchmark.

See training_summary.json for checkpoint provenance and SHA256SUMS for file hashes.

Credits and license

Base model and original ComfyUI package: well9472/Nanosaur2-670M, MIT. This package derives from upstream revision ff8e403f40cb091bb57ef5a72a2a1e962da24aa0; original file hashes are recorded in upstream.json.

model.py and nodes.py include the inpainting integration, and inpaint.py contains the new adapter implementation. Replacing those files with upstream copies would remove this support. The included upstream train_lora.py trains base-model LoRAs; it is not the training script for this inpainting adapter.

LLLite inspiration: kohya-ss/Anima-LLLite. Adapter and integration by levzalt. Released under the MIT license.