SANA-Denoiser - Our Architecture Design¶

Repurposing SANA 1.6B DiT as an image restoration model. Combines efficient linear attention with Paired Training for Restoration and Temporal Tiling via Block Causal Linear Attention.

Why SANA for Restoration¶

SANA is 10x smaller, 4x fewer tokens, linear complexity. For restoration where we need high-res processing, this is decisive.

Architecture Changes (Minimal)¶

1. Input Conditioning: Channel Concat¶

degraded → DC-AE.encode → condition_latents [B, 32, H, W]
target   → DC-AE.encode → latents           [B, 32, H, W]

x_t = (1-σ)*noise + σ*latents               [B, 32, H, W]
model_input = concat([x_t, condition_latents], dim=1)  [B, 64, H, W]

projection = Conv2d(64, 32, 1)               # 1x1 conv, ~1K params
# Identity init for noise channels, zero init for condition channels
# At step 0: model = pretrained T2I behavior
# Condition signal learned gradually during fine-tuning

model(projection(model_input), timestep, text_embeddings)

Total new parameters: 1,024 (32 x 32 x 1 x 1 conv kernel). Compare: ControlNet = ~800M.

2. Text Conditioning for Degradation Type¶

Prompt describes what to restore: - "Remove gaussian noise, restore sharp details" - "Remove JPEG compression artifacts" - "Enhance this low-light image" - "Clean and restore this image"

Leverages SANA's Gemma-2-2B text encoder for degradation-type understanding.

3. Temporal Tiling for High-Resolution¶

For images > training resolution (e.g., 4K product photos):

4096x4096 image
  ↓ split into overlapping 1024px tiles (raster scan)
  ↓ each tile: DC-AE encode → 32x32x32 latent
  ↓ denoise with Block Causal Linear Attention
  ↓    (running sum S, Z from previous tiles = global context)
  ↓ stitch latents with linear blending in overlap
  ↓ DC-AE decode full stitched latent
4096x4096 restored image

Memory: constant O(D^2) cache + one tile latent. Processes any resolution.

Training Strategy¶

Phase 1: LoRA (fast iteration)¶

• Rank 32, target: attn.to_q/k/v/out + input projection conv
• 512px, 10K steps, DIV2K + Flickr2K synthetic degradation
• Evaluate: does it learn to denoise at all?

Phase 2: Full Fine-Tune (if LoRA insufficient)¶

• Unfreeze all transformer params + projection
• VAE stays frozen
• Gradient checkpointing for memory
• Curriculum: 512px → 1024px

Phase 3: Temporal Tiling (inference-only first)¶

• No retraining needed - causal attention is native to linear attention
• Just implement the tile loop + S, Z accumulation
• If quality insufficient: fine-tune with multi-tile samples

Dataset¶

Source: DIV2K (800) + Flickr2K (2650) = 3450 clean images Degradation: 5-8 variants per image = 17K-28K pairs

Evaluation Targets¶

Project Files¶

happyin-research/
├── sana-fm/
│   ├── data/paired_dataset.py      ← paired loader
│   ├── data/degradation.py         ← degradation functions
│   ├── configs/img2img_denoise.yaml
│   └── train_flowmatching.py       ← modified compute_loss
├── sana-denoiser/
│   ├── prepare_dataset.py          ← DIV2K + Flickr2K + degradations
│   ├── train.py                    ← wrapper
│   ├── temporal_tiling.py          ← tile-as-sequence inference
│   └── eval/benchmark.py           ← vs SwinIR, NAFNet

Risk Assessment¶