Qwen VAE DeGrid NAFNet 1x

The short normie explanation:

This is a 1x "upscale" (image restoration) model. It is designed to remove the subtle Qwen/WAN VAE 'grid artifact' that appears in many images outputted by diffusion models trained with those VAEs (such as Anima). It will not work to it's full potential in ComfyUI without a custom node - Please see the usage section below (Reason: I chose an arch which doesn't work as automagically as others. No I'm not gonna retrain it. You can do that)

For those that want to indulge in some information about the model:

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 (for example, this introduces color changes and halos/sharpening artifacts, but it does remove the grid, too)

• 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.
  
  Image showing what ComfyUI nodes are needed for non-custom node usage (CivitAI would error when I tried to upload image directly here 🙂🙂🙂🙂):
  https://cdn-uploads.huggingface.co/production/uploads/650f8271f2e0c7dc7ede5517/jIDvOiKvrYnTOWF8noo0G.png

• 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 on this CivitAI model card (showcase images) but also on the HuggingFace repository's assets folder. It would be best to download them yourself 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

HF repo: https://huggingface.co/DraconicDragon/NAFNet-VAE-DeGrid

Training code (i think configs too):

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

Btw model works fine with https://huggingface.co/Anzhc/Qwen2D-Anime-VAE decoder finetune of Qwen VAE too.

FAQ that no one ever asked

Q: Why?

A: Didn't know of any solution to this. Nyquist Notch shader helped but it's effect is too weak.

Q: Support for sd webui forge/A1111/reforge/reforge class neo/minecraft forge/sd.next/ersatzForge/stability matrix?
A: Too much effort for me, and I don't use those programs anyway/anymore. But ComfyUI nodes and training code is there for code reference to shovel into LLM. Maybe it works out of the box, but I doubt it.

Q: Will there be future updates?
A: Unlikely, I see this model as finished. The most likely update in any way would be that I retrain this model on a potentially more easy to use architecture, but even that is extremly unlikely because this model is already pretty good for what it is IMO. There's a finetune by Anzhc which improves on my v1.1 by a very tiny bit, you might not see any difference, but you can dream well knowing that the numbers are slightly better, and personally, I also believe it is perceptually a very tiny about better - as is reflected by the numbers.