SnapOtter/docs/specs/2026-04-15-restore-photo-design.md

Restore-Photo: State-of-the-Art Redesign

Date: 2026-04-15
Status: Approved
Scope: Full rewrite of the restore-photo pipeline and settings UI

---

Overview

Replace the current restore-photo pipeline with a single, always-best-quality pipeline backed by state-of-the-art open-source models running fully locally inside the Docker container. No presets, no quality tiers — the pipeline always runs at maximum quality. GPU is used when available; CPU is always a supported fallback. Users control which steps run via toggles, not how well.

---

Models

All models are ONNX (or TFLite for BlazeFace) and downloaded at Docker build time via hf_hub_download. If a model file is missing at runtime, the affected step is skipped gracefully with a logged warning.

Model	Purpose	Format	Approx. Size	HuggingFace Source
lama_fp32.onnx	Inpainting fallback (when SD model absent)	ONNX fp32	~200 MB	Carve/LaMa-ONNX
codeformer.onnx	Face restoration (primary)	ONNX	~170 MB	facefusion/models-3.0.0
gfpgan_1.4.onnx	Face restoration fallback	ONNX	~330 MB	facefusion/models-3.0.0
nafnet_denoise.onnx	Neural denoising	ONNX	~68 MB	Export from megvii-research/NAFNet official weights (NAFNet-SIDD-width32)
ddcolor.onnx	Colorization (shared with colorize tool)	ONNX	~450 MB	existing
ct2_colorization.onnx	Higher-quality colorization	ONNX	~850 MB	Export from shuchenweng/CT2 (ICCV 2023) official weights
realesrgan_x2.onnx	Super-resolution 2× (shared with upscale tool)	ONNX	~65 MB	existing
realesrgan_x4.onnx	Super-resolution 4× (shared with upscale tool)	ONNX	~65 MB	existing
mediapipe_blazeface.tflite	Face detection	TFLite float16	~2 MB	Google MediaPipe (existing)
sd_inpaint_fp16.onnx	Diffusion inpainting (primary)	ONNX fp16	~1.7 GB	runwayml/stable-diffusion-inpainting quantized

Total new model footprint: ~3.9 GB. Models shared with other tools are not re-downloaded.

---

Processing Pipeline

Stages run sequentially. Each stage is a self-contained function: takes image + config, returns modified image. GPU provider priority: CUDAExecutionProvider → CoreMLExecutionProvider → CPUExecutionProvider.

Input image
    │
    ▼
[1] Pre-processing
    - EXIF auto-orient
    - Detect grayscale (is_grayscale flag) — mean channel diff < 5.0
    - Compute adaptive tile size from image resolution
    │
    ▼
[2] Damage Detection  (always runs to produce mask for step 3)
    - 8-angle morphological line kernels (0°, 22.5°, 45°, 67.5°, 90°, 112.5°, 135°, 157.5°)
    - Laplacian-based stain/fade/water-damage detection layer (new)
    - Combines top-hat + black-hat transforms at full resolution
    - Output: damage_mask at full resolution
    │
    ▼
[3] Inpainting  (if scratchRemoval=true AND damage coverage ≥ 0.1%)
    Primary: SD-inpainting ONNX fp16
        - Overlapping tile-stitch (tile=512, stride=384)
        - Poisson-blend seam correction
    Fallback (if SD model absent): LaMa tiled
        - tile=768, stride=512, feathered Gaussian blend
    │
    ▼
[4] Face Enhancement  (if faceEnhancement=true)
    - MediaPipe BlazeFace detection (unchanged)
    - Skip faces < 48px
    - Crop with 80% expansion for context
    - CodeFormer at 512×512 (primary)
    - GFPGAN at 512×512 as fallback if CodeFormer output SSIM vs. original face crop < 0.6
    - Paste-back with face-aware feathered elliptical mask + color-transfer correction
    │
    ▼
[5] Denoising  (if denoise=true)
    - NAFNet ONNX, tiled at full resolution
    - denoiseStrength (0–100) maps to NAFNet sigma parameter
    - Falls back to NLMeans if NAFNet model absent
    │
    ▼
[6] Super Resolution  (if superResolution=true)
    - Real-ESRGAN x4 if input < 2 MP
    - Real-ESRGAN x2 if input ≥ 2 MP
    - Tiled inference with seam blending (reuse upscale tool implementation)
    │
    ▼
[7] Colorization  (if colorize=true)
    - CT2 tiled (tile=512, stride=384)
    - Falls back to DDColor tiled if CT2 model absent
    - colorizeStrength (0–100) controls LAB a/b blend intensity
    - Preserves original luminance channel
    │
    ▼
Output image

Progress milestones (stderr JSON {"progress": N, "stage": "..."}):

%	Stage
5	Pre-processing
10	Damage detection
25	Inpainting
45	Face enhancement
60	Denoising
75	Super resolution
90	Colorization
95	Saving output
100	Complete (set by API route)

---

API Changes

Removed

• mode: "auto" | "light" | "heavy" — removed entirely

Unchanged

• scratchRemoval: boolean (default true)
• faceEnhancement: boolean (default true)
• fidelity: number 0–1 (default 0.7) — CodeFormer fidelity
• denoise: boolean (default true)
• denoiseStrength: number 0–100 (default 40)
• colorize: boolean (default false)

Added

• superResolution: boolean (default true) — allow user to disable SR if they want original dimensions
• colorizeStrength: number 0–100 (default 85) — was hardcoded 0.85 internally

Timeout

Bridge timeout raised from 300s → 1200s to accommodate SD-inpainting on CPU for large images.

---

Frontend Changes (restore-photo-settings.tsx)

• Remove 3-column mode button group (auto / light / heavy)
• Add Super Resolution checkbox (default on)
• Add Colorize Strength slider (0–100, step 5) — shown when colorize is on, same pattern as fidelity slider
• denoiseStrength slider remains; now controls NAFNet sigma rather than NLMeans h

No other UI changes. Settings panel stays compact.

---

i18n Additions (packages/shared/src/i18n/en.ts)

superResolution: "Super Resolution",
colorizeStrength: "Colorize Strength",

Remove: restoreMode, restoreModeAuto, restoreModeLight, restoreModeHeavy (if present).

---

Error Handling

Scenario	Behavior
SD model missing	Fall back to LaMa tiled; log warning
CT2 model missing	Fall back to DDColor tiled; log warning
NAFNet model missing	Fall back to NLMeans; log warning
GFPGAN model missing	Skip GFPGAN fallback; CodeFormer only
No damage detected (< 0.1% coverage)	Skip inpainting entirely
Tiling dimension mismatch	Fall back to 512px single-pass for that stage; log warning
CUDA unavailable	Silently use CPU provider
SD-inpainting CPU timeout	Bridge timeout is 1200s; if exceeded, API returns 504 with clear message

---

Testing

Unit (Vitest)

• API Zod schema accepts superResolution and colorizeStrength
• API schema rejects old mode field
• colorizeStrength default is 85

Integration (Vitest)

• POST restore-photo with fixture image → assert output dimensions ≥ input (SR ran)
• POST with superResolution: false → assert output dimensions equal input
• POST with colorize: true → assert output is not identical to input

Python (pytest)

• Each pipeline stage tested in isolation with small fixture image:
  • Damage detection returns a binary mask with same H×W as input
  • LaMa tiled returns same-shape image
  • NAFNet tiled returns same-shape image
  • Real-ESRGAN x2 returns 2× dimensions
  • CT2 tiled returns same-shape image

E2E (Playwright)

• Upload fixture old photo → verify before-after slider renders with non-blank result
• Verify Super Resolution checkbox appears in settings panel

---

File Paths Affected

File	Change
packages/ai/python/restore.py	Full rewrite of pipeline
packages/ai/python/download_models.py	Add NAFNet, CT2, SD-inpainting, GFPGAN download entries
packages/ai/src/restoration.ts	Update RestoreOptions interface
apps/api/src/routes/tools/restore-photo.ts	Update Zod schema, raise timeout
apps/web/src/components/tools/restore-photo-settings.tsx	Remove mode group, add SR + colorizeStrength controls
packages/shared/src/i18n/en.ts	Add/remove i18n keys
packages/shared/src/constants.ts	Update tool timeout constant if applicable