SnJake
/

JPG_Noise_Remover

@@ -7,98 +7,74 @@ tags:
   - image-restoration
   - image-enhancement
   - denoising
   - comfyui
   - pytorch
 ---
-# About this project
-This project is a personal experiment created out of curiosity. The main part of the code was generated by an AI assistant, and my task was to set the goal, prepare the data, run the training and evaluate the result. The model is trained to remove artifacts from images (JPEG, noise) and even shows good results.
-## Artifacts Remover UNet
-This is a lightweight UNet-based model trained to remove JPEG compression artifacts and additive Gaussian noise from images. The model is ideal for integration into image processing pipelines, including the popular ComfyUI framework.
 ## Examples
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/658814fd586088fd274d8cc1/ipja2Z2qzV4-39SB_6zDa.png)
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/658814fd586088fd274d8cc1/I4kPu6EYMhBpID0zFxZYm.png)
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/658814fd586088fd274d8cc1/BijSCHP_f-GBQ2a4BNUJw.png)
-### How to use in ComfyUI
- **Step 1: Install the Node**
-1.  Open a terminal or command prompt.
-2.  Navigate to your ComfyUI `custom_nodes` directory.
-    ```bash
-    # Example for Windows
-    cd D:\ComfyUI\custom_nodes\
-    # Example for Linux
-    cd ~/ComfyUI/custom_nodes/
-    ```
-3.  Clone this repository into the `custom_nodes` folder:
-    ```bash
-    git clone https://github.com/SnJake/SnJake_JPG_Artifacts_Noise_Cleaner.git
-    ```
-4.  **Install Dependencies**: Now, you need to install the required Python packages. The command depends on which version of ComfyUI you are using.
-    *   **For standard ComfyUI installations (with venv):**
-        1.  Make sure your ComfyUI virtual environment (`venv`) is activated.
-        2.  Navigate into the new node directory and install the requirements:
-            ```bash
-            cd SnJake_JPG_Artifacts_Noise_Cleaner
-            pip install -r requirements.txt
-            ```
-    *   **For Portable ComfyUI installations:**
-        1.  Navigate back to the **root** of your portable ComfyUI directory (e.g., `D:\ComfyUI_windows_portable`).
-        2.  Run the following command to use the embedded Python to install the requirements. *Do not activate any venv.*
-            ```bash
-            python_embeded\python.exe -m pip install -r ComfyUI\custom_nodes\SnJake_JPG_Artifacts_Noise_Cleaner\requirements.txt
-            ```
- **Step 2: Install the Model Weights (Or you can skip this step, the node will download the model weights itself after starting the Queue)**
-1.  Navigate to your `ComfyUI/models/` directory.
-2.  Create a folder named `artifacts_remover` inside it, if it doesn't already exist.
-3.  Download the model weights file (`.pt` or `.safetensors`).
-4.  Place the downloaded weights file into the `ComfyUI/models/artifacts_remover/` directory.
-### Step 3: Restart
-Restart ComfyUI completely. The new node will be available in the "Add Node" menu.
-### Training Details
-The model was trained on a dataset of approximately 30,000 high-quality images, primarily consisting of anime-style art. Instead of using pre-degraded images, the training process generated (degraded, clean) image pairs on-the-fly.
-*   **Architecture**: The network is a `UNetRestorer` built with `ResidualBlock`s for deep feature extraction. To enhance important features, the deeper levels of the encoder utilize the Convolutional Block Attention Module (CBAM). The model employs a final residual connection, learning to predict the difference (`clean - degraded`) rather than the entire clean image.
-*   **Degradation Process**: Each clean image patch was subjected to a sequence of randomly ordered degradations:
-    *   **JPEG Compression**: A random quality level was chosen between 5 and 85.
-    *   **Gaussian Noise**: Gaussian noise was added with a standard deviation randomly selected from the range [0.0, 7.0].
-    *   **Identity Mapping**: With a 20% probability (`--clean-prob 0.2`), the input image was left clean (not degraded). This encourages the model to preserve details when no artifacts are present.
-*   **Training Procedure**:
-    *   **Optimizer**: AdamW with a learning rate of `2e-4` and weight decay of `1e-4`.
-    *   **Learning Rate Scheduler**: A Cosine Annealing scheduler with a linear warmup phase of 2000 steps was used.
-    *   **Batch & Patch Size**: The model was trained with a batch size of 12 using 320x320 pixel patches.
-    *   **Loss Function**: A comprehensive, multi-component loss function was employed to balance pixel accuracy, structural integrity, and perceptual quality:
-        *   **Primary Loss**: A weighted sum of `0.7 * CharbonnierLoss` (a smooth L1 variant) and `0.3 * MixL1SSIM`. The `MixL1SSIM` component itself was weighted with `alpha=0.9`, combining L1 loss and a structural similarity term (`0.9*L1 + 0.1*(1-SSIM)`).
-        *   **Edge Loss**: `GradientLoss` was added with a weight of 0.15 (`--edge-loss-w 0.15`) to penalize blurry edges and promote sharpness.
-        *   **High-Frequency Error Norm (HFEN)**: To better preserve fine textures and details, `HFENLoss` was included with a weight of 0.12 (`--hfen-w 0.12`).
-        *   **Identity Loss**: For the 20% of samples where the input was clean, an additional L1 loss with a weight of 0.5 (`--id-loss-w 0.5`) was calculated between the model's output and the input. This forces the network to act as an identity function for high-quality images, preventing it from introducing blur or altering details.
-    *   **Techniques**: Training was accelerated using Automatic Mixed Precision (AMP) with the `bfloat16` data type. An Exponential Moving Average (EMA) of the model's weights (`decay=0.999`) was maintained to produce a more stable and generalized final model for inference.
-### Limitations and Potential Issues
-*   The model was trained on primarily consisting of anime-style art. Results on photos, line art, or text may be suboptimal.
-*   With very high levels of noise or artifacts beyond the training range, the model may hallucinate details or over-smooth the image.
-*   The model might interpret very fine, low-contrast textures (e.g., fabric, sand) as noise and smooth them out. For such cases, use the `blend` parameter in the node to mix back some of the original detail.
-*   The model does not correct for other types of degradation, such as motion blur, chromatic aberrations, or optical flaws.

   - image-restoration
   - image-enhancement
   - denoising
+  - deblurring
   - comfyui
   - pytorch
+  - nafnet
 ---
+# JPG Artifacts & Noise Remover (V2)
+This is a lightweight model designed to remove JPEG compression artifacts, digital noise, and slight blur from images. Ideally suited for anime/illustration art, but works reasonably well on photos.
+**Version 2 Update:**
+*   **New Architecture:** Switched to a NAFNet-based UNet (State-of-the-Art blocks for restoration).
+*   **Larger Dataset:** Trained on ~40,000 high-quality images from Danbooru2024.
+*   **Improved Training:** Added Perceptual Loss, HFEN, and blur degradation handling.
+> **Which version to pick?** V2 generally restores better, but it is also cautious and may decide an image is already clean. If V2 feels too conservative for a specific image, try the V1 weights (`best_ema_15E.safetensors`) and keep the result you like more.
 ## Examples
+![1_Ready](https://cdn-uploads.huggingface.co/production/uploads/658814fd586088fd274d8cc1/7xw2xkYC1dTSyArMz9Tbn.png)
+![2_Ready](https://cdn-uploads.huggingface.co/production/uploads/658814fd586088fd274d8cc1/xQmp5wZBBDJquSSTEbzJ2.png)
+![3_Ready](https://cdn-uploads.huggingface.co/production/uploads/658814fd586088fd274d8cc1/LAvBqBvTIecZz3FEAclzp.png)
+## How to use in ComfyUI
+The model is designed to work with the **[JPG & Noise Remover ComfyUI Node](https://github.com/SnJake/SnJake_JPG_Artifacts_Noise_Cleaner)**.
+1.  **Install the Node:** Follow instructions in the [GitHub Repository](https://github.com/SnJake/SnJake_JPG_Artifacts_Noise_Cleaner).
+2.  **Download Weights:** Download the `.pt` or `.safetensors` v2 file from this repository.
+3.  **Place Weights:** Put the file in `ComfyUI/models/artifacts_remover/`.
+4.  **Select Model:** Select the new weight file in the node settings. Ensure `base_ch` is set to **64**.
+## Training Details (V2)
+The goal was to create a restorer that not only removes noise but preserves fine structural details without "plastic" smoothing.
+### Dataset
+Trained on **40,000 images** from the [Danbooru2024](https://huggingface.co/datasets/deepghs/danbooru2024) dataset (anime/illustration style).
+### Architecture: NAFNet-based UNet
+The model uses a U-Net structure but replaces standard convolutional blocks with **NAFBlocks** (Nonlinear Activation Free Blocks).
+*   **SimpleGate:** Replaces complex activation functions with element-wise multiplication.
+*   **LayerNorm2d:** Stabilizes training.
+*   **Simplified Channel Attention (SCA):** captures global context efficiently.
+*   **Base Channels:** 64
+### Degradation Pipeline
+The model is trained on "on-the-fly" generated pairs. The degradation pipeline is more aggressive than V1:
+1.  **Blur:** Random downscale-upscale (probability 50%, scale down to 0.85x) to simulate soft blur.
+2.  **JPEG Compression:** Quality 10 - 85.
+3.  **Gaussian Noise:** Standard deviation 0 - 20.0.
+4.  **Identity:** 2% of images are left clean to teach the model not to alter good images.
+### Loss Function
+A composite loss function was used to balance pixel accuracy and perceptual quality:
+*   **Pixel Loss:** Charbonnier Loss (0.7) + MixL1SSIM (0.3).
+*   **Perceptual Loss (VGG19):** Weight 0.05. Helps generating realistic textures.
+*   **HFEN (High-Frequency Error Norm):** Weight 0.1. Enforces edge reconstruction.
+*   **Gradient/Edge Loss:** Weight 0.2.
+*   **Identity Loss:** Weight 0.02 (applied on clean images).
+### Training Config
+*   **Optimizer:** AdamW (`lr=2e-4`, `wd=1e-4`).
+*   **Scheduler:** Cosine Annealing (2000 warmup steps).
+*   **Precision:** BFloat16 (AMP).
+*   **Patch Size:** 288x288.
+*   **Batch Size:** 4 (Accumulated to effective 8).
+## Limitations
+*   Primarily trained on anime/2D art.
+*   May struggle with extremely heavy motion blur (since it trained mostly on slight downscale blur).