NanoSD: Edge Efficient Foundation Model for Real Time Image Restoration

TL;DR

NanoSD proposes an edge-friendly diffusion backbone by hardware-aware reengineering of the SD 1.5 U-Net, block-wise distillation, and end-to-end VAE refinement to yield a Pareto frontier of backbones balancing accuracy, latency, and size. It employs a hardware-aware block decomposition, feature-wise distillation, and Bayesian optimization with taFID to preserve SD 1.5's generative prior while dramatically reducing on-device latency. The resulting NanoSD family delivers real-time on-device image restoration across SR, FR, deblurring, dehazing, deraining, desnowing, and monocular depth estimation, demonstrated on SM8750 and Apple ANE with tile-based processing. This approach generalizes across platforms, showing that latency-aware architecture search coupled with distillation can unlock practical diffusion-based restoration on edge devices while preserving the teacher model's priors.

Abstract

Latent diffusion models such as Stable Diffusion 1.5 offer strong generative priors that are highly valuable for image restoration, yet their full pipelines remain too computationally heavy for deployment on edge devices. Existing lightweight variants predominantly compress the denoising U-Net or reduce the diffusion trajectory, which disrupts the underlying latent manifold and limits generalization beyond a single task. We introduce NanoSD, a family of Pareto-optimal diffusion foundation models distilled from Stable Diffusion 1.5 through network surgery, feature-wise generative distillation, and structured architectural scaling jointly applied to the U-Net and the VAE encoder-decoder. This full-pipeline co-design preserves the generative prior while producing models that occupy distinct operating points along the accuracy-latency-size frontier (e.g., 130M-315M parameters, achieving real-time inference down to 20ms on mobile-class NPUs). We show that parameter reduction alone does not correlate with hardware efficiency, and we provide an analysis revealing how architectural balance, feature routing, and latent-space preservation jointly shape true on-device latency. When used as a drop-in backbone, NanoSD enables state-of-the-art performance across image super-resolution, image deblurring, face restoration, and monocular depth estimation, outperforming prior lightweight diffusion models in both perceptual quality and practical deployability. NanoSD establishes a general-purpose diffusion foundation model family suitable for real-time visual generation and restoration on edge devices.

Figures (15)

• Figure 1: We present NanoSD, an edge-efficient diffusion model designed for image restoration applications. Its core principle is to harness the extensive visual priors embedded in contemporary generative image models while maintaining computational feasibility for edge deployment. NanoSD demonstrates practical applicability across diverse low-level vision tasks in real-world scenarios.
• Figure 2: Overview of the NanoSD framework. (a) Baseline SD 1.5 U– Net architecture, shown with skip connections removed for readability. (b) Hardware-aware search space construction: for each of the six retained stages (three encoders and three decoders), we derive shape-preserving residual/attention variants that are profiled on the target edge device for latency and parameter cost. (c) Feature-wise generative distillation: each candidate block is distilled independently from its corresponding SD 1.5 teacher block using an $\ell_2$ feature-matching loss. (d) Combinatorial assembly and evaluation: a decision vector specifies one distilled block per stage, producing a structurally valid U– Net; each assembled model is evaluated on taFID and either on-device latency or parameter count via Bayesian optimization. (e) Latency--taFID Pareto frontier obtained from the search; green points denote Pareto-optimal models, and non-Pareto points are shown for comparison. (f) Parameter--taFID Pareto frontier. Both plots also include a hand-tuned baseline and the Segmind TinySD model, which fall far from the frontier. (g) The resulting NanoSD family: seven Pareto-optimal architectures spanning different accuracy--efficiency trade-offs (bubble size inversely proportional to accuracy). (h) Final selected architecture, NanoSD (Model 2), used for all downstream experiments; skip connections are omitted for clarity.
• Figure 3: Generated images using edge efficient NanoSD.
• Figure 4: Qualitative comparisons of different SR methods. Please zoom in for a better view. Refer to Supplemental for more results.
• Figure 5: Visual comparisons of various face restoration methods. Please zoom in for a better view. Refer to Supplemental for more results.