From a98c7bc1ac21cefd670893f81165a4a87007cedb Mon Sep 17 00:00:00 2001
From: Andrii Ryzhkov <andrii.ryzhkov@pm.me>
Date: Wed, 15 Apr 2026 11:59:22 +0200
Subject: [PATCH] Enable NAFNet model

---
 README.md                        | 1 +
 darktable_ai/config.py           | 2 ++
 darktable_ai/convert.py          | 3 +++
 models/denoise-nafnet/.skip      | 0
 models/denoise-nafnet/model.yaml | 7 +++++--
 pyproject.toml                   | 9 +++++++++
 6 files changed, 20 insertions(+), 2 deletions(-)
 delete mode 100644 models/denoise-nafnet/.skip

diff --git a/README.md b/README.md
index f99275b..e9f4853 100644
--- a/README.md
+++ b/README.md
@@ -8,6 +8,7 @@ Currently targets the ONNX backend. The pipeline is designed to support addition
 
 | Model                                                                         | Task    | Description                                     |
 |-------------------------------------------------------------------------------|---------|-------------------------------------------------|
+| [`denoise-nafnet`](models/denoise-nafnet/README.md)                           | denoise | NAFNet denoiser trained on SIDD dataset         |
 | [`denoise-nind`](models/denoise-nind/README.md)                               | denoise | UNet denoiser trained on NIND dataset           |
 | [`embed-openclip-vitb32`](models/embed-openclip-vitb32/README.md)             | embed   | OpenCLIP ViT-B/32 text/image embeddings         |
 | [`mask-object-sam21-base-plus`](models/mask-object-sam21-base-plus/README.md) | mask    | SAM 2.1 Hiera Base Plus for interactive masking |
diff --git a/darktable_ai/config.py b/darktable_ai/config.py
index e659957..3cccddf 100644
--- a/darktable_ai/config.py
+++ b/darktable_ai/config.py
@@ -45,6 +45,7 @@ class ModelConfig:
     skip: bool = False
 
     model_card: dict[str, str] = field(default_factory=dict)
+    attributes: dict = field(default_factory=dict)
 
     repo: RepoConfig | None = None
     checkpoints: list[Checkpoint] = field(default_factory=list)
@@ -118,6 +119,7 @@ def load_model_config(model_dir: Path, root_dir: Path) -> ModelConfig:
         tiling=data.get("tiling", False),
         dep_group=data.get("dep_group", "core"),
         model_card=data.get("model_card", {}),
+        attributes=data.get("attributes", {}),
         skip=skip,
         repo=repo,
         checkpoints=checkpoints,
diff --git a/darktable_ai/convert.py b/darktable_ai/convert.py
index 8966382..97d1241 100644
--- a/darktable_ai/convert.py
+++ b/darktable_ai/convert.py
@@ -49,6 +49,9 @@ def generate_config_json(config: ModelConfig) -> None:
     if config.model_card:
         data["model_card"] = config.model_card
 
+    if config.attributes:
+        data["attributes"] = config.attributes
+
     config_file.write_text(json.dumps(data, indent=4, ensure_ascii=False) + "\n")
     print(f"  Generated: {config_file}")
 
diff --git a/models/denoise-nafnet/.skip b/models/denoise-nafnet/.skip
deleted file mode 100644
index e69de29..0000000
diff --git a/models/denoise-nafnet/model.yaml b/models/denoise-nafnet/model.yaml
index 196c6bd..3d1c584 100644
--- a/models/denoise-nafnet/model.yaml
+++ b/models/denoise-nafnet/model.yaml
@@ -3,10 +3,14 @@ name: "denoise nafnet small"
 description: "NAFNet denoiser trained on SIDD dataset"
 task: denoise
 version: "1.0"
+arch: nafnet
 tiling: true
 type: single
 dep_group: nafnet
 
+attributes:
+  shadow_boost: true
+
 model_card:
   long_description: "NAFNet (Nonlinear Activation Free Network) lightweight denoiser trained on the SIDD smartphone denoising dataset"
   scope: "single-image denoising"
@@ -20,7 +24,6 @@ model_card:
 
 repo:
   submodule: vendor/NAFNet
-  setup: "python setup.py develop --no_cuda_ext"
 
 checkpoints:
   - url: "https://drive.google.com/file/d/1lsByk21Xw-6aW7epCwOQxvm6HYCQZPHZ/view?usp=sharing"
@@ -33,4 +36,4 @@ convert:
       checkpoint: "{temp}/NAFNet-SIDD-width32.pth"
       output: "{output}/model.onnx"
       opset: 17
-      fp16: true
+      fp16: false
diff --git a/pyproject.toml b/pyproject.toml
index 560b2bd..bfcdef6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,6 +26,15 @@ core = [
 nafnet = [
     {include-group = "core"},
     "onnxconverter-common",
+    "addict",
+    "lmdb",
+    "scikit-image",
+    "scipy",
+    "tqdm",
+    "yapf",
+    "requests",
+    "future",
+    "tb-nightly",
 ]
 nind = [
     {include-group = "core"},