ArtisanLabs/artisan-upscaler

Artisan Upscaler

GitHub | Hugging Face

4x single-image super-resolution based on DAT (Dual Aggregation Transformer), trained on ~4M high-resolution images.

Model

Architecture	DAT D512 — 8 residual groups, multi-scale U-body, neighborhood attention
Parameters	89.8M
Scale	4x
Attention	Neighborhood attention (15x15 kernel, dual dilation [1, 3]) via PyTorch flex_attention (or NATTEN if installed)
Training	1.2M steps, progressive 256->512 crops, GAN fine-tuning from 700K
Precision	BF16 weights (safetensors)

Results

Each strip shows: LR input (left) | Bicubic 4x | Artisan Upscaler 4x | Ground Truth

Photography

Art & Portraits

Quick Start

pip install torch torchvision safetensors timm huggingface_hub

Download weights from Hugging Face:

huggingface-cli download ArtisanLabs/artisan-upscaler weights/artisan_upscaler_bf16.safetensors --local-dir .

import torch
from config import DATConfig
from DAT import DAT
from safetensors.torch import load_file

# Load model
cfg = DATConfig(attn_type="natten", multiscale=True, hr_refine_blocks=2)
model = DAT(
    img_size=cfg.img_size, in_chans=cfg.in_chans, embed_dim=cfg.embed_dim,
    split_size=cfg.split_size, depth=cfg.depth, num_heads=cfg.num_heads,
    expansion_factor=cfg.expansion_factor, qkv_bias=cfg.qkv_bias,
    drop_path_rate=cfg.drop_path_rate, upscale=cfg.upscale,
    img_range=cfg.img_range, upsampler=cfg.upsampler,
    resi_connection=cfg.resi_connection, attn_type=cfg.attn_type,
    natten_kernel=cfg.natten_kernel, natten_dilation=cfg.natten_dilation,
    use_chk=cfg.use_chk, multiscale=cfg.multiscale,
    ms_enc_groups=cfg.ms_enc_groups, ms_dec_groups=cfg.ms_dec_groups,
    hr_refine_blocks=cfg.hr_refine_blocks,
)
sd = load_file("weights/artisan_upscaler_bf16.safetensors", device="cpu")
model.load_state_dict(sd, strict=False)
model = model.to("cuda", dtype=torch.bfloat16, memory_format=torch.channels_last).eval()
model = torch.compile(model)  # optional, ~2x faster after warmup

# Upscale
lr = torch.randn(1, 3, 128, 128, device="cuda", dtype=torch.bfloat16)
lr = lr.to(memory_format=torch.channels_last)
with torch.no_grad():
    sr = model(lr)  # (1, 3, 512, 512)

CLI Usage

# Single image
python inference.py weights/artisan_upscaler_bf16.safetensors input.png -o output.png

# Directory of images
python inference.py weights/artisan_upscaler_bf16.safetensors input_dir/ -o output_dir/

# Smaller tiles for low-VRAM GPUs
python inference.py weights/artisan_upscaler_bf16.safetensors input.png --tile 64

# torch.compile for faster throughput (slower first image)
python inference.py weights/artisan_upscaler_bf16.safetensors input.png --compile

Options

Flag	Default	Description
--tile	128	LR tile size (128 = 512px HR tiles)
--overlap	32	Overlap between tiles for seamless blending
--precision	bf16	bf16 or fp32
--compile	off	torch.compile for ~2x faster per-image throughput
--device	cuda	Device (cuda, cuda:0, cpu)

VRAM Usage

Measured with torch.compile, BF16:

Tile Size	VRAM	Notes
64	~0.3 GB	Fits on any GPU
128 (default)	~0.7 GB	Recommended
256	~3 GB	Fewer tiles, faster for large images

Images larger than the tile size are automatically split into overlapping tiles and blended with a Hann window for seamless output.

Attention Backend

The model uses neighborhood attention (15x15 local window). Two backends are supported:

1. PyTorch flex_attention (default) — works out of the box with PyTorch >= 2.5. Use torch.compile for best performance.
2. NATTEN — fused CUDA kernels, faster on Ampere/Hopper/Blackwell. Auto-detected if installed.

No configuration needed — the model picks the best available backend automatically.

Architecture

Input (H, W, 3)
  |
  Conv 3x3 -> shallow features (H, W, 512)
  |
  Multi-scale U-body:
    Encoder: 3 residual groups at full res
    Downsample 2x
    Bottleneck: 2 residual groups at half res
    Upsample 2x
    Decoder: 3 residual groups at full res
  |
  Each residual group: 3 DATB blocks + 1x1 conv
  Each DATB: LayerNorm -> NA/Channel Attn -> LayerNorm -> SGFN
  |
  Conv 3x3 -> residual connection
  |
  PixelShuffle 4x -> HR refinement (2 ResBlocks)
  |
Output (4H, 4W, 3)

Training

Trained on ~4.1M images from 10 datasets:

• PD12M (3.8M, 3-8 megapixel filtered)
• LAION-HR (106K, clean + watermarked tracks)
• FFHQ (70K faces)
• Art Portraits (64K curated)
• LSDIR (41K diverse)
• DIV2K, Flickr2K, HQ-50K, MS ImagePairs

Degradation: random bicubic/bilinear/area downsampling with optional JPEG compression. GAN fine-tuning with hinge loss + projected discriminator from step 700K.

Files

inference.py       — CLI inference (tiled, any resolution)
config.py          — DATConfig dataclass
DAT.py             — Model architecture
vectorized_ops.py  — Windowing and masking utilities
weights/           — Model checkpoints (safetensors)

Dependencies

• PyTorch >= 2.5
• torchvision
• safetensors
• timm

Optional: NATTEN for faster attention kernels.

License

Apache 2.0

Citation

@software{artisan_upscaler,
  title={Artisan Upscaler: 4x Super-Resolution with Dual Aggregation Transformer},
  author={TheArtisanAI},
  year={2025},
  url={https://github.com/TheArtisanAI/artisan-upscaler}
}