DraconicDragon/NAFNet-VAE-DeGrid

Info

A quick and dirty NAFNet-small based image restoration model for Qwen/Wan VAE grid artifact removal/mitigation, trained with 0 experience and a lot of AI handholding.

• Aim: Something on par with or hopefully better than Nyquist Notch shader - Remove the grid without changing/deforming/destroying anything
• Limitations: 1. Maybe realism, idk; 2. This model is only trained on what the VAE changes and very likely won't touch any extra noise or such your diffusion model generated

Finetune (v1.1) - I trained a small 30k step finetune with slightly different settings and a smaller dataset where I filtered out most film grain-heavy images.
It seems to perform a somewhat noticeable amount better; And also doesn't suffer from a repetitive pattern that occurs on artifact-heavy images.

---

Usage

The model outputs a residual rather than the final image, so the residual needs to be added back to the original image.

ComfyUI

Custom nodes are not required - but you need custom nodes linked below to use the model to it's full potential.

• No custom nodes: You can use the model only with ComfyUI's built-in nodes by loading it with Load Upscale Model, passing it through Upscale Image (using Model), and then using Image Blend with a Float node set to -1.0 to add the residual back to the original image. However, ComfyUI clamps the model output to the normal image range, so this mainly gives you the "Dark Pixels Mainly" behavior (see below). The positive residuals are preserved while the negative residuals are clamped away. This result is visually similar to Nyquist Notch shader.

• Custom nodes:

  For full control over the residual, use the ComfyUI NAFNet Residual node pack. It provides three modes:

  • Full - apply the entire residual, removing both dark and bright artifacts.
  • Dark Pixels Mainly - apply only positive residuals. This leaves the bright corrections out and is the same result as with ComfyUI core nodes above. (Also similar to Nyquist Notch shader's result)
  • Bright Pixels Mainly - apply only negative residuals.

  The node pack also includes a NAFNet Residual Patch method, which patches ComfyUI to let you use the normal Load Upscale Model → Upscale Image (using Model) workflow while applying the residual automagically. The patch can use the same three modes, as well as Disabled to return to normal ComfyUI/Spandrel behavior, in case of any issues.

  See the node pack README for installation and more details.

---

Comparison

You can find images for one comparison in the assets folder. Download them and put them in your favourite image comparer or something.
Areas to look out for are:

• Dark/bright end of tail on the left
• Leaves in the background, easier to see on left
• Base of the cat ears, hair too
• The purple fuzzy wrist cuffs
• Left and right parts of the skirt
• Eyes
• Shadows on most images in general will have subtle artifacts

---

Some training info

The model was trained in FP16 with automatic mixed precision for 100k steps on my good old 3060 12GB, which took about 19-20 hours at a patch size of 256px and batch size of 8.

Dataset: The dataset consists mostly of image pairs of paintings, illustrations, and the like, with a focus on anime-oriented artworks. I don't know how it performs on realism.

Dataset image pairs:

• GT/Ground Truth: PNG images
• Low quality/corrupted: GT images passed through a Qwen VAE encode-decode pass

While the model was trained on the artifacts of outputs from Qwen VAE, it seems to work with Wan 2.1 VAE outputs too, even if the outputs are slightly different.

Note that the model might perform slightly worse since it wasn't explicitly trained with Wan 2.1 VAE artifacts.

---

The model cannot imitate/replicate film grain (and averages/smooths/blurs it), not that it needs to be able to reproduce it, but it might struggle with generated images that would have film grain-like texture if not for the VAE grid artifacts. Maybe I will revisit it because the dataset could definitely need many improvements which would help. The finetune fixes most of that struggle and a repetitive pattern likely caused by that tiny extra punch of FFT loss that was there, or the large film grain heavy portion of the full dataset

Checkpoints

You can find PyTorch checkpoints as they came from training in the train/ subfolder:

• 100k steps - pretrain
• 30k steps - finetune

Training code (i think configs too):

https://github.com/DraconicDragon/NAFNet-c