release model and provide dependencies

DISCLAIMER

@@ -0,0 +1 @@
+This project is maintained independently and is not affiliated with any organization. Commercial use requires prior written permission.

GRDFNet.py

@@ -0,0 +1,132 @@
+import torch
+from torch import nn
+
+
+def conv3x3(in_channels, out_channels, bias=True):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias
+    )
+
+
+class LWGRB(nn.Module):
+    def __init__(self, channels: int, bias: bool = True, identity=True):
+        super().__init__()
+        self.identity = identity
+        self.conv1 = conv3x3(channels, channels, bias)
+        self.act = nn.LeakyReLU(0.1, inplace=True)
+        self.conv2 = conv3x3(channels, channels, bias)
+        nn.init.zeros_(self.conv2.weight)
+        nn.init.zeros_(self.conv2.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r = self.conv2(self.act(self.conv1(x)))
+        a = torch.sigmoid(r)
+        return x + r * a
+
+
+class LWDRB(nn.Module):
+    def __init__(self, c, bias=True, dil=4):
+        super().__init__()
+        self.conv1 = conv3x3(c, c, bias)
+        self.act = nn.LeakyReLU(0.1, inplace=True)
+        self.conv2 = nn.Conv2d(c, c, 3, padding=dil, dilation=dil, bias=bias)
+        nn.init.zeros_(self.conv2.weight)
+        nn.init.zeros_(self.conv2.bias)
+
+    def forward(self, x):
+        r = self.conv2(self.act(self.conv1(x)))
+        return x + r
+
+
+class LWGRBShuffle(nn.Module):
+    def __init__(self, in_ch, out_ch, scale, bias: bool = True):
+        super().__init__()
+        self.expand = conv3x3(in_ch, out_ch * scale * scale, bias)
+        self.up = nn.PixelShuffle(scale)
+        self.refine = LWGRB(out_ch, bias=bias)
+
+    def forward(self, x):
+        x = self.expand(x)
+        x = self.up(x)
+        return self.refine(x)
+
+
+class GRDFNet(nn.Module):
+    """
+    GRDFNet (Gated + Residual Dilated Fast Network)
+    Configurable block stacking via integer `num_sets`.
+
+    Structure:
+      Stem: 3xLWGRB + 2xLWDRB
+      Repeated: num_sets x [LWGRB + LWDRB]
+    """
+
+    def __init__(
+        self,
+        num_in_ch: int = 3,
+        num_out_ch: int = 3,
+        feature_channels: int = 32,
+        upscale: int = 1,
+        bias: bool = True,
+        norm: bool = False,
+        img_range: float = 1.0,
+        rgb_mean=(0.5, 0.5, 0.5),
+        num_sets: int = 3,
+    ):
+        super().__init__()
+
+        self.in_ch = num_in_ch
+        self.out_ch = num_out_ch
+        self.c = feature_channels
+        self.scale = upscale
+        self.img_range = img_range
+        self.gamma = nn.Parameter(torch.tensor(0.5))
+        self.num_sets = num_sets
+
+        self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
+        if not norm:
+            self.register_buffer("no_norm", torch.zeros(1))
+        else:
+            self.no_norm = None
+
+        self.head = conv3x3(self.in_ch, self.c, bias)
+        self.body = self._make_body(bias)
+        self.tail = conv3x3(self.c, self.out_ch, bias)
+
+        if self.scale == 1:
+            self.upsample0 = nn.Identity()
+        else:
+            self.upsample0 = LWGRBShuffle(self.out_ch, self.out_ch, self.scale)
+
+    def _make_body(self, bias: bool):
+        blocks = [
+            LWGRB(self.c, bias),
+            LWGRB(self.c, bias),
+            LWGRB(self.c, bias),
+            LWDRB(self.c, bias),
+            LWDRB(self.c, bias),
+        ]
+        for _ in range(self.num_sets):
+            blocks += [LWGRB(self.c, bias), LWDRB(self.c, bias)]
+        return nn.Sequential(*blocks)
+
+    @property
+    def is_norm(self) -> bool:
+        return getattr(self, "no_norm", None) is None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.is_norm:
+            self.mean = self.mean.type_as(x)
+            x = (x - self.mean) * self.img_range
+
+        feat = self.head(x)
+        feat = self.body(feat)
+        out_feat = self.tail(feat)
+        out_feat = (1.0 - self.gamma) * x + self.gamma * out_feat
+
+        out = self.upsample0(out_feat)
+
+        if self.is_norm:
+            out = out / self.img_range + self.mean
+
+        return out

LICENSE

@@ -0,0 +1,25 @@
+MIT License (With Commercial Use Restriction)
+
+Copyright (c) 2025 Nicholas Lane
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+1. The above copyright notice and this permission notice (including the
+commercial use restriction below) shall be included in all copies or
+substantial portions of the Software.
+2. Commercial use of the Software, including any derivative works, is
+prohibited unless prior explicit written permission is granted by the
+copyright holder.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

NEOSR/__init__.py

@@ -0,0 +1,268 @@
+import importlib
+import os
+import random
+from copy import deepcopy
+from functools import partial
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.utils.data
+from torch.utils import data
+from torch.utils.data.sampler import Sampler
+
+from neosr.utils import get_root_logger, scandir
+from neosr.utils.dist_util import get_dist_info
+from neosr.utils.registry import DATASET_REGISTRY
+
+__all__ = ["build_dataloader", "build_dataset"]
+
+# automatically scan and import dataset modules for registry
+# scan all the files under the data folder with '_dataset' in file names
+data_folder = Path(Path(__file__).resolve()).parent
+dataset_filenames = [
+    Path(Path(v).name).stem
+    for v in scandir(str(data_folder))
+    if v.endswith("_dataset.py")
+]
+
+def build_dataset(dataset_opt: dict[str, Any]):
+    """Build dataset from options.
+
+    Args:
+    ----
+        dataset_opt (dict): Configuration for dataset. It must contain:
+            type (str): Dataset type.
+
+    """
+    dataset_opt = deepcopy(dataset_opt)
+    dataset_type = dataset_opt["type"]
+    dataset_cls = DATASET_REGISTRY.get(dataset_type)  # type: ignore[assignment]
+    logger = get_root_logger()
+
+    def _maybe_cache_to_ram() -> None:
+        default_cache = str(dataset_type).lower().startswith("cached")
+        cache_requested = dataset_opt.get("cache_in_ram", default_cache)
+        dataroot_lq = dataset_opt.get("dataroot_lq")
+        dataroot_gt = dataset_opt.get("dataroot_gt")
+        if not cache_requested:
+            logger.debug(
+                "Skipping RAM cache preload for dataset '%s' (cache_in_ram disabled).",
+                dataset_type,
+            )
+            return
+        if not dataroot_lq or not dataroot_gt:
+            logger.warning(
+                "cache_in_ram enabled for dataset '%s' but dataroots are missing. Skipping preload.",
+                dataset_type,
+            )
+            return
+
+        lq_root = Path(dataroot_lq)
+        gt_root = Path(dataroot_gt)
+        if not lq_root.exists() or not gt_root.exists():
+            logger.warning(
+                "cache_in_ram enabled for dataset '%s' but dataroots do not exist. "
+                "LQ: %s, GT: %s. Skipping preload.",
+                dataset_type,
+                dataroot_lq,
+                dataroot_gt,
+            )
+            return
+
+        try:
+            from cv2 import IMREAD_UNCHANGED, imread  # type: ignore[attr-defined]
+        except ImportError as exc:  # pragma: no cover - OpenCV should be available
+            logger.warning(
+                "cache_in_ram enabled for dataset '%s' but OpenCV is unavailable: %s",
+                dataset_type,
+                exc,
+            )
+            return
+
+        def collect_files(root: Path) -> dict[str, Path]:
+            return {
+                str(path.relative_to(root)).replace("\\", "/"): path
+                for path in root.rglob("*")
+                if path.is_file()
+            }
+
+        lq_files = collect_files(lq_root)
+        gt_files = collect_files(gt_root)
+        common_rel_paths = sorted(set(lq_files) & set(gt_files))
+        if not common_rel_paths:
+            logger.warning(
+                "cache_in_ram enabled for dataset '%s' but no matching LQ/GT pairs were found.",
+                dataset_type,
+            )
+            return
+
+        if len(common_rel_paths) != len(lq_files) or len(common_rel_paths) != len(gt_files):
+            logger.info(
+                "cache_in_ram: restricting to %d matched pairs (dropped %d LQ, %d GT).",
+                len(common_rel_paths),
+                len(lq_files) - len(common_rel_paths),
+                len(gt_files) - len(common_rel_paths),
+            )
+
+        lq_images: list[np.ndarray] = []
+        gt_images: list[np.ndarray] = []
+        total_bytes = 0
+        for rel_path in common_rel_paths:
+            lq_img = imread(str(lq_files[rel_path]), IMREAD_UNCHANGED)
+            gt_img = imread(str(gt_files[rel_path]), IMREAD_UNCHANGED)
+            if lq_img is None or gt_img is None:
+                logger.warning(
+                    "cache_in_ram: failed to read pair '%s'. Aborting preload.",
+                    rel_path,
+                )
+                lq_images.clear()
+                gt_images.clear()
+                break
+            lq_img = np.ascontiguousarray(lq_img).astype(np.float32) / 255.0
+            gt_img = np.ascontiguousarray(gt_img).astype(np.float32) / 255.0
+            lq_images.append(lq_img)
+            gt_images.append(gt_img)
+            total_bytes += lq_img.nbytes + gt_img.nbytes
+        if not lq_images or not gt_images:
+            return
+
+        dataset_opt["_cached_lq_images"] = lq_images
+        dataset_opt["_cached_gt_images"] = gt_images
+        dataset_opt["_cached_rel_paths"] = common_rel_paths
+        dataset_opt["_cached_lq_paths"] = [str(lq_files[p]) for p in common_rel_paths]
+        dataset_opt["_cached_gt_paths"] = [str(gt_files[p]) for p in common_rel_paths]
+        dataset_opt.setdefault("scale", dataset_opt.get("scale", 1))
+        logger.info(
+            "Preloaded %d image pairs for dataset '%s' into RAM (approx %.2f GiB).",
+            len(common_rel_paths),
+            dataset_type,
+            total_bytes / (1024**3),
+        )
+
+    _maybe_cache_to_ram()
+    dataset = dataset_cls(dataset_opt)  # type: ignore[operator]
+    cache_active = getattr(dataset, "cache_in_ram_active", False) or bool(
+        dataset_opt.get("_cached_lq_images")
+    )
+    if cache_active:
+        setattr(dataset, "cache_in_ram_active", True)
+    logger.info(f"Dataset [{dataset.__class__.__name__}] is built.")
+    return dataset
+
+
+def worker_init_fn(worker_id: int, num_workers: int, rank: int, seed: int) -> None:
+    # Set the worker seed to num_workers * rank + worker_id + seed
+    worker_seed = num_workers * rank + worker_id + seed
+    # NOTE: set seed on old generator as a precaution, but
+    # it is redundand since we use np.random.Generator
+    np.random.seed(worker_seed)  # noqa: NPY002
+    torch.manual_seed(worker_seed)
+    random.seed(worker_seed)
+
+
+def build_dataloader(
+    dataset: data.Dataset,
+    dataset_opt: dict[str, Any],
+    num_gpu: int = 1,
+    dist: bool = False,
+    sampler: Sampler | None = None,
+    seed: int | None = None,
+) -> data.DataLoader:
+    """Build dataloader.
+
+    Args:
+    ----
+        dataset (torch.utils.data.Dataset): Dataset.
+        dataset_opt (dict): Dataset options. It contains the following keys:
+            phase (str): 'train' or 'val'.
+            num_worker_per_gpu (int): Number of workers for each GPU.
+            batch_size (int): Training batch size for each GPU.
+        num_gpu (int): Number of GPUs. Used only in the train phase.
+            Default: 1.
+        dist (bool): Whether in distributed training. Used only in the train
+            phase. Default: False.
+        sampler (torch.utils.data.sampler): Data sampler. Default: None.
+        seed (int | None): Seed. Default: None
+
+    """
+    phase = dataset_opt["phase"]
+    rank, _ = get_dist_info()
+    logger = get_root_logger()
+
+    # train
+    if phase == "train":
+        cpu_count = os.cpu_count() or 1
+        workers_per_gpu_opt = dataset_opt.get("num_worker_per_gpu")
+        # Heuristic tuned for high-end hosts: keep workers high enough to saturate I/O.
+        if workers_per_gpu_opt is None or workers_per_gpu_opt == "auto":
+            gpu_slots = num_gpu if num_gpu > 0 else 1
+            workers_per_gpu = max(4, min(32, cpu_count // gpu_slots))
+        else:
+            workers_per_gpu = int(workers_per_gpu_opt)
+        workers_per_gpu = max(workers_per_gpu, 1)
+        num_workers = workers_per_gpu
+
+        if dist:  # distributed training
+            batch_size = dataset_opt["batch_size"]
+        else:  # non-distributed training
+            multiplier = 1 if num_gpu == 0 else num_gpu
+            batch_size = dataset_opt["batch_size"] * multiplier
+            num_workers *= multiplier
+
+        if "prefetch_factor" in dataset_opt:
+            prefetch_factor = dataset_opt["prefetch_factor"]
+        else:
+            # High worker counts benefit from deeper prefetch queues.
+            prefetch_factor = max(4, min(16, workers_per_gpu))
+
+        cache_in_ram_active = getattr(dataset, "cache_in_ram_active", False)
+        if os.name == "nt" and cache_in_ram_active and num_workers > 0:
+            logger.warning(
+                "Windows detected with RAM-cached dataset '%s'; forcing num_workers from %d to 0 "
+                "to avoid spawn pickling overhead.",
+                dataset.__class__.__name__,
+                num_workers,
+            )
+            num_workers = 0
+
+        dataloader_args = {
+            "dataset": dataset,
+            "batch_size": batch_size,
+            "shuffle": False,
+            "num_workers": num_workers,
+            "sampler": sampler,
+            "drop_last": True,
+        }
+        if num_workers > 0:
+            dataloader_args["prefetch_factor"] = prefetch_factor
+        if sampler is None:
+            dataloader_args["shuffle"] = True
+        dataloader_args["worker_init_fn"] = (
+            partial(worker_init_fn, num_workers=num_workers, rank=rank, seed=seed)
+            if seed is not None
+            else None
+        )
+
+    # val
+    elif phase in {"val", "test"}:
+        dataloader_args = {
+            "dataset": dataset,
+            "batch_size": 1,
+            "shuffle": False,
+            "num_workers": 0,
+        }
+    else:
+        msg = f"Wrong dataset phase: {phase}. Supported ones are 'train', 'val' and 'test'."
+        raise ValueError(msg)
+
+    dataloader_args["pin_memory"] = dataset_opt.get("pin_memory", True)
+    if "persistent_workers" in dataset_opt:
+        dataloader_args["persistent_workers"] = dataset_opt["persistent_workers"]
+    else:
+        dataloader_args["persistent_workers"] = dataloader_args.get("num_workers", 0) > 0
+        if dataloader_args["persistent_workers"] and dataloader_args["num_workers"] == 0:
+            dataloader_args["persistent_workers"] = False
+
+    return data.DataLoader(**dataloader_args)

NEOSR/cachedsr_dataset.py

@@ -0,0 +1,157 @@
+from pathlib import Path
+from typing import Any
+
+import cv2
+import numpy as np
+from torch.utils import data
+from torchvision.transforms.functional import normalize
+
+from neosr.data.transforms import basic_augment, paired_random_crop
+from neosr.utils import get_root_logger, img2tensor
+from neosr.utils.registry import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class CachedSRDataset(data.Dataset):
+    """Paired SR dataset that keeps image pairs in host RAM for fast sampling."""
+
+    def __init__(self, opt: dict[str, Any]) -> None:
+        super().__init__()
+        self.opt = opt
+        self.logger = get_root_logger()
+        self.scale = opt.get("scale", 1)
+        self.patch_size = opt.get("patch_size")
+        self.phase = opt.get("phase", "train")
+        self.mean = opt.get("mean")
+        self.std = opt.get("std")
+        self.color = opt.get("color", None) != "y"
+        self.use_hflip = opt.get("use_hflip", True)
+        self.use_rot = opt.get("use_rot", True)
+
+        cached_lq = opt.get("_cached_lq_images")
+        cached_gt = opt.get("_cached_gt_images")
+        cached_lq_paths = opt.get("_cached_lq_paths")
+        cached_gt_paths = opt.get("_cached_gt_paths")
+        cached_rel = opt.get("_cached_rel_paths")
+
+        self.cache_in_ram_active = False
+        if cached_lq and cached_gt:
+            self.lq_images = cached_lq
+            self.gt_images = cached_gt
+            self.lq_paths = cached_lq_paths or [
+                f"{opt.get('dataroot_lq', 'lq')}::{idx}" for idx in range(len(self.lq_images))
+            ]
+            self.gt_paths = cached_gt_paths or [
+                f"{opt.get('dataroot_gt', 'gt')}::{idx}" for idx in range(len(self.gt_images))
+            ]
+            self.rel_paths = cached_rel or [str(idx) for idx in range(len(self.lq_images))]
+            self.cache_in_ram_active = True
+            self.logger.info(
+                "CachedSRDataset reusing %d preloaded pairs from build_dataset cache.",
+                len(self.lq_images),
+            )
+        else:
+            self.logger.info("CachedSRDataset preloading images locally (no shared cache provided).")
+            dataroot_lq = opt.get("dataroot_lq")
+            dataroot_gt = opt.get("dataroot_gt")
+            if not dataroot_lq or not dataroot_gt:
+                msg = "CachedSRDataset requires dataroot_lq and dataroot_gt when cache is absent."
+                raise ValueError(msg)
+            self._load_from_disk(dataroot_lq, dataroot_gt)
+            self.cache_in_ram_active = True
+
+        assert len(self.lq_images) == len(self.gt_images), "Unmatched LQ/GT counts."
+        self.length = len(self.lq_images)
+
+    def _load_from_disk(self, dataroot_lq: str, dataroot_gt: str) -> None:
+        lq_root = Path(dataroot_lq)
+        gt_root = Path(dataroot_gt)
+        if not lq_root.exists() or not gt_root.exists():
+            msg = f"CachedSRDataset dataroots do not exist. LQ: {dataroot_lq}, GT: {dataroot_gt}"
+            raise FileNotFoundError(msg)
+
+        def collect(root: Path) -> dict[str, Path]:
+            return {
+                str(path.relative_to(root)).replace("\\", "/"): path
+                for path in root.rglob("*")
+                if path.is_file()
+            }
+
+        lq_files = collect(lq_root)
+        gt_files = collect(gt_root)
+        keys = sorted(set(lq_files) & set(gt_files))
+        if not keys:
+            msg = "CachedSRDataset could not find any matched LQ/GT pairs."
+            raise RuntimeError(msg)
+        dropped_lq = len(lq_files) - len(keys)
+        dropped_gt = len(gt_files) - len(keys)
+        if dropped_lq or dropped_gt:
+            self.logger.warning(
+                "CachedSRDataset dropping unmatched pairs (LQ: -%d, GT: -%d).",
+                dropped_lq,
+                dropped_gt,
+            )
+
+        self.lq_images = []
+        self.gt_images = []
+        self.lq_paths = []
+        self.gt_paths = []
+        self.rel_paths = keys
+
+        for rel in keys:
+            lq_img = cv2.imread(str(lq_files[rel]), cv2.IMREAD_UNCHANGED)
+            gt_img = cv2.imread(str(gt_files[rel]), cv2.IMREAD_UNCHANGED)
+            if lq_img is None or gt_img is None:
+                msg = f"Failed to read cached pair {rel}."
+                raise RuntimeError(msg)
+            lq_img = np.ascontiguousarray(lq_img).astype(np.float32) / 255.0
+            gt_img = np.ascontiguousarray(gt_img).astype(np.float32) / 255.0
+            self.lq_images.append(lq_img)
+            self.gt_images.append(gt_img)
+            self.lq_paths.append(str(lq_files[rel]))
+            self.gt_paths.append(str(gt_files[rel]))
+
+    def __len__(self) -> int:
+        return self.length
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        lq_img = self.lq_images[index]
+        gt_img = self.gt_images[index]
+        gt_path = self.gt_paths[index] if index < len(self.gt_paths) else None
+        lq_path = self.lq_paths[index] if index < len(self.lq_paths) else None
+
+        # Training branch: random crop + augment
+        if self.phase == "train" and self.patch_size:
+            gt_img, lq_img = paired_random_crop(
+                gt_img, lq_img, self.patch_size, self.scale, gt_path or self.rel_paths[index]
+            )
+            gt_img = np.ascontiguousarray(gt_img)
+            lq_img = np.ascontiguousarray(lq_img)
+            gt_img, lq_img = basic_augment(
+                [gt_img, lq_img],
+                hflip=self.use_hflip,
+                rotation=self.use_rot,
+            )  # type: ignore[reportAssignmentType]
+            gt_img = np.ascontiguousarray(gt_img)
+            lq_img = np.ascontiguousarray(lq_img)
+        else:
+            # Align GT to LQ size when evaluating
+            h_lq, w_lq = lq_img.shape[:2]
+            gt_img = gt_img[: h_lq * self.scale, : w_lq * self.scale, ...]
+            gt_img = np.ascontiguousarray(gt_img)
+            lq_img = np.ascontiguousarray(lq_img)
+
+        gt_tensor, lq_tensor = img2tensor(
+            [gt_img, lq_img], bgr2rgb=True, float32=True, color=self.color
+        )
+
+        if self.mean is not None or self.std is not None:
+            normalize(lq_tensor, self.mean, self.std, inplace=True)  # type: ignore[arg-type]
+            normalize(gt_tensor, self.mean, self.std, inplace=True)  # type: ignore[arg-type]
+
+        return {
+            "lq": lq_tensor,
+            "gt": gt_tensor,
+            "lq_path": lq_path or self.rel_paths[index],
+            "gt_path": gt_path or self.rel_paths[index],
+        }

NEOSR/notes.md

@@ -0,0 +1,2 @@
+Add cachedsr_dataset and update `__init__.py` in NEOSR for faster training if sytsem RAM allows.
+Minimal training TOML provided for non-GAN training.

NEOSR/train_grdfnet.toml

@@ -0,0 +1,78 @@
+
+name = 'NAME-Here'
+model_type = 'image'
+scale = 2 
+use_amp = true
+bfloat16 = true
+fast_matmul = true
+[datasets.train]
+type = 'CachedSRDataset'
+dataroot_gt = 'GT IMAGES' 
+dataroot_lq = 'LQ IMAGES'
+patch_size = 96
+batch_size = 32
+accumulate = 1
+augmentation = [ 'resizemix', 'cutblur']    # 'mixup', 'cutmix', 'resizemix', 'cutblur'
+aug_prob = [ 0.25, 0.15 ]
+use_hflip = true
+use_rot = true
+
+
+[datasets.val]
+name = 'val'
+type = 'paired'
+dataroot_gt = 'GT IMAGES' 
+dataroot_lq = 'LQ IMAGES'
+
+[val]
+val_freq = 5000
+[val.metrics.psnr]
+type = 'calculate_psnr'
+
+[train]
+ema = 0.995
+match_lq_colors = false
+
+[network_g]
+type = 'GRDFNet'
+num_sets = 3
+feature_channels = 32
+
+[train.scheduler]
+type = 'cosineannealing'
+T_max = 150000
+eta_min = 1e-9 
+
+[train.optim_g]
+type = 'adam'
+lr = 1e-4                    
+betas = [0.9, 0.999]
+weight_decay = 5e-7
+
+[train.pixel_opt]
+type = 'HuberLoss'
+loss_weight = 1.0
+reduction = 'mean'
+
+[train.gw_opt]
+type = "gw_loss"
+loss_weight = 0.15
+criterion = "chc_loss"
+corner = true
+
+[train.consistency_opt]
+type = 'consistency_loss'
+loss_weight = 0.25
+criterion = 'chc' # 'l1'
+blur = true
+cosim = true
+saturation = 1.0
+brightness = 1.0
+
+
+[logger]
+total_iter = 150000
+save_checkpoint_freq = 5000
+use_tb_logger = true
+#save_tb_img = true
+print_freq = 500

README.md

@@ -1,14 +1,40 @@
 ---
 title: GRDFNet
-emoji: 🚀
 colorFrom: red
 colorTo: yellow
 sdk: gradio
 sdk_version: 5.44.1
 app_file: app.py
-pinned: false
-license: other
+license: mit
 short_description: A lightweight image restoration architecture
+pinned: true
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# GRDFNet
+
+GRDFNet is a lightweight image restoration network that combines gated and dilated residual blocks to deliver strong perceptual quality with modest compute requirements.
+
+## Recommended Configurations
+
+- `num_sets = 3`, `feature_channels = 32`: strong quality while staying fast for most desktop workloads.
+- `num_sets = 6`, `feature_channels = 48`: highest quality configuration; expect roughly a 4x slowdown versus the 32-channel model.
+- `num_sets = 3`, `feature_channels = 24`: suggested for lightly compressed video inference; typically 50~75% faster than the 32-channel variant when deployed with TensorRT.
+
+## Performance Snapshot
+
+Example TensorRT run on an NVIDIA RTX 4080 Super (16 GB):
+```
+DEBUG: TensorRT initialized. Setting shape.
+DEBUG: Shape set. Getting output shape.
+[INFO] Input: 1280x720 -> 1280x720 -> ModelOut: 1280x720 @ 30000/1001 fps
+DEBUG: Before NVENC initialization.
+[prof] frames=209 avg=208.2 fps
+[prof] frames=431 avg=215.1 fps
+[prof] frames=654 avg=217.3 fps
+[INFO] Processed 709 frames in 3.280s -> 216.2 FPS
+```
+
+## Resources
+
+- Model weights: https://huggingface.co/nicholasLane/GRDFNet
+- Hosted demo: https://huggingface.co/spaces/nicholasLane/GRDFNet

SPANDREL/GRDFNet.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+import torch
+from torch import nn as nn
+
+from spandrel.util import store_hyperparameters
+
+def conv3x3(in_channels, out_channels, bias=True):
+    return nn.Conv2d(
+        in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias
+    )
+
+
+class LWGRB(nn.Module):
+    def __init__(self, channels: int, bias: bool = True, identity=True):
+        super().__init__()
+        self.identity = identity
+        self.conv1 = conv3x3(channels, channels, bias)
+        self.act = nn.LeakyReLU(0.1, inplace=True)
+        self.conv2 = conv3x3(channels, channels, bias)
+        nn.init.zeros_(self.conv2.weight)
+        nn.init.zeros_(self.conv2.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r = self.conv2(self.act(self.conv1(x)))
+        a = torch.sigmoid(r)
+        return x + r * a
+
+
+class LWDRB(nn.Module):
+    def __init__(self, c, bias=True, dil=4):
+        super().__init__()
+        self.conv1 = conv3x3(c, c, bias)
+        self.act = nn.LeakyReLU(0.1, inplace=True)
+        self.conv2 = nn.Conv2d(c, c, 3, padding=dil, dilation=dil, bias=bias)
+        nn.init.zeros_(self.conv2.weight)
+        nn.init.zeros_(self.conv2.bias)
+
+    def forward(self, x):
+        r = self.conv2(self.act(self.conv1(x)))
+        return x + r
+
+
+class LWGRBShuffle(nn.Module):
+    def __init__(self, in_ch, out_ch, scale, bias: bool = True):
+        super().__init__()
+        self.expand = conv3x3(in_ch, out_ch * scale * scale, bias)
+        self.up = nn.PixelShuffle(scale)
+        self.refine = LWGRB(out_ch, bias=bias)
+
+    def forward(self, x):
+        x = self.expand(x)
+        x = self.up(x)
+        return self.refine(x)
+
+
+@store_hyperparameters()
+class GRDFNet(nn.Module):
+    """
+    GRDFNet (Gated + Residual Dilated Fast Network)
+    Configurable block stacking via integer `num_sets`.
+
+    Structure:
+      Stem: 3xLWGRB + 2xLWDRB
+      Repeated: num_sets x [LWGRB + LWDRB]
+    """
+
+    def __init__(
+        self,
+        num_in_ch: int = 3,
+        num_out_ch: int = 3,
+        feature_channels: int = 32,
+        upscale: int = 1,
+        bias: bool = True,
+        norm: bool = False,
+        img_range: float = 1.0,
+        rgb_mean=(0.5, 0.5, 0.5),
+        num_sets: int = 3,
+    ):
+        super().__init__()
+
+        self.in_ch = num_in_ch
+        self.out_ch = num_out_ch
+        self.c = feature_channels
+        self.scale = upscale
+        self.img_range = img_range
+        self.gamma = nn.Parameter(torch.tensor(0.5))
+        self.num_sets = num_sets
+
+        self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
+        if not norm:
+            self.register_buffer("no_norm", torch.zeros(1))
+        else:
+            self.no_norm = None
+
+        self.head = conv3x3(self.in_ch, self.c, bias)
+        self.body = self._make_body(bias)
+        self.tail = conv3x3(self.c, self.out_ch, bias)
+
+        if self.scale == 1:
+            self.upsample0 = nn.Identity()
+        else:
+            self.upsample0 = LWGRBShuffle(self.out_ch, self.out_ch, self.scale)
+
+    def _make_body(self, bias: bool):
+        blocks = [
+            LWGRB(self.c, bias),
+            LWGRB(self.c, bias),
+            LWGRB(self.c, bias),
+            LWDRB(self.c, bias),
+            LWDRB(self.c, bias),
+        ]
+        for _ in range(self.num_sets):
+            blocks += [LWGRB(self.c, bias), LWDRB(self.c, bias)]
+        return nn.Sequential(*blocks)
+
+    @property
+    def is_norm(self) -> bool:
+        return getattr(self, "no_norm", None) is None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.is_norm:
+            self.mean = self.mean.type_as(x)
+            x = (x - self.mean) * self.img_range
+
+        feat = self.head(x)
+        feat = self.body(feat)
+        out_feat = self.tail(feat)
+        out_feat = (1.0 - self.gamma) * x + self.gamma * out_feat
+
+        out = self.upsample0(out_feat)
+
+        if self.is_norm:
+            out = out / self.img_range + self.mean
+
+        return out

SPANDREL/__init__.py

@@ -0,0 +1,111 @@
+import math
+from typing import Union
+import torch
+from typing_extensions import override
+
+from spandrel.util import KeyCondition
+
+from ...__helpers.model_descriptor import Architecture, ImageModelDescriptor, StateDict
+from .arch.GRDFNet import GRDFNet
+
+
+def _infer_scale(state_dict):
+    try:
+        expand_weight = state_dict["upsample0.expand.weight"]
+        in_ch = expand_weight.shape[1]
+        ratio = expand_weight.shape[0] // in_ch
+        scale = int(math.isqrt(ratio))
+        assert (
+            scale * scale == ratio
+        ), "Unexpected expand weight shape"  # 1, 2, 4, 8, 16...
+        print(f"scale: {scale}")
+    except:
+        scale = 1
+
+    return scale
+
+
+def _infer_num_sets(state_dict: dict[str, torch.Tensor]) -> int:
+    body_keys = [k for k in state_dict.keys() if k.startswith("body.")]
+    block_indices = {k.split(".")[1] for k in body_keys if k.split(".")[1].isdigit()}
+    num_blocks = len(block_indices)
+
+    if num_blocks <= 5:
+        # minimal model (fallback)
+        return 0
+
+    # each set contributes +2 blocks beyond the 5-block stem
+    num_sets = max((num_blocks - 5) // 2, 0)
+    return num_sets
+
+
+def compute_receptive_fields(num_sets: int) -> tuple[int, int]:
+    # 3xLWGRB + 2xLWDRB + (num_sets x [1 LWGRB + 1 LWDRB])
+    n_grb = 3 + num_sets
+    n_drb = 2 + num_sets
+    trf = 1 + 2 + n_grb * (2 * 1) * 2 + n_drb * (2 * 4) * 2 + 2
+    erf = int(trf * 0.7)
+    return erf, trf
+
+
+class GRDFNetArch(Architecture[GRDFNet]):
+    def __init__(self) -> None:
+        super().__init__(
+            id="GRDFNet",
+            detect=KeyCondition.has_any(
+                "head.weight",
+                "tail.weight",
+            ),
+        )
+
+    @override
+    def load(self, state_dict: StateDict) -> ImageModelDescriptor[GRDFNet]:
+        num_in_ch: int = 3
+        num_out_ch: int = 3
+        feature_channels: int = 32
+        norm = True
+        img_range = 255.0
+        rgb_mean = (0.4488, 0.4371, 0.4040)
+
+        if "no_norm" in state_dict:
+            norm = False
+            state_dict["no_norm"] = torch.zeros(1)
+
+        feature_channels, num_in_ch, _, _ = state_dict["head.weight"].shape
+        num_out_ch, _, _, _ = state_dict["tail.weight"].shape
+
+        scale = _infer_scale(state_dict)
+        num_sets = _infer_num_sets(state_dict)
+        state_params = sum(v.numel() for v in state_dict.values())
+        erf, trf = compute_receptive_fields(num_sets)
+
+        model = GRDFNet(
+            num_in_ch=num_in_ch,
+            num_out_ch=num_out_ch,
+            feature_channels=feature_channels,
+            upscale=scale,
+            norm=norm,
+            img_range=img_range,
+            rgb_mean=rgb_mean,
+            num_sets=num_sets,
+        )
+
+        return ImageModelDescriptor(
+            model,
+            state_dict,
+            architecture=self,
+            purpose="SR",
+            tags=[
+                f"{feature_channels}ch",
+                f"x{scale}",
+                f"nS:{num_sets}",
+                f"Params:{state_params}",
+                f"ERF:{erf}",
+                f"TRF:{trf}",
+            ],
+            supports_half=True,
+            supports_bfloat16=True,
+            scale=scale,
+            input_channels=num_in_ch,
+            output_channels=num_out_ch,
+        )

SPANDREL/notes.md

@@ -0,0 +1,4 @@
+Ensure you add GRDFNet to the main registry
+```
+ArchSupport.from_architecture(GRDFNet.GRDFNetArch())
+```

app.py

@@ -0,0 +1,20 @@
+# app.py
+import os, types, gradio as gr
+
+app_code = os.environ.get("APP_CODE", "")
+
+
+def run_hidden(code_str):
+    module = types.ModuleType("hidden_app")
+    exec(code_str, module.__dict__)
+    if hasattr(module, "main"):
+        return module.main()
+    else:
+        raise RuntimeError("Secret app must define main()")
+
+
+if app_code:
+    demo = run_hidden(app_code)
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+else:
+    raise RuntimeError("APP_CODE is not set. Please add it in your Space Secrets.")