Enable NAFNet denoise model

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 |

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,

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}")
 

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

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"},