KumaPower/CSA

ICCV 2019 · Modern PyTorch implementation

Coherent Semantic Attention for Image Inpainting

Recover missing image regions with a rough prediction and context-aware refinement.

Hongyu Liu · Bin Jiang · Yi Xiao · Chao Yang

Paper · Code · Quickstart · Training · Datasets

_{Stage 2 · step 97,397. Columns: masked input · restored · original.
Selected examples.}

The author's CSA checkpoint combines a separately trained rough U-Net with CSA refinement. This modern release uses random masks, L1 + LPIPS, and the original image and feature PatchGANs. It updates the original paper's training recipe; see implementation notes.

The released refinement checkpoint is step 97,397, with a frozen rough parent at step 156,181. It was selected by visual review, rather than automatic best-metric selection.

How it works

Train rough prediction first, then freeze it and train refinement. CSA operates on 32 × 32 encoder features at the 256 × 256 training resolution. This diagram follows the visual layout of Figure 2 in the paper and shows the modern training recipe. Editable SVG.

Try the checkpoint

1. Install

Verified Linux/NVIDIA environment: Python 3.10.14 · PyTorch 2.6.0 · torchvision 0.21.0 · CUDA 12.4. The official wheel pairing needs a compatible NVIDIA driver, with no separate CUDA toolkit installation. The wheel below installs CSA without a Git checkout.

python3.10 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch==2.6.0 torchvision==0.21.0 \
  --index-url https://download.pytorch.org/whl/cu124
python -m pip install \
  "https://huggingface.co/KumaPower/CSA/resolve/main/csa_inpainting-0.3.0-py3-none-any.whl"
python - <<'PY'
import torch, torchvision
assert torch.cuda.is_available(), "CUDA is unavailable"
print(torch.__version__, torchvision.__version__)
PY

For CPU, macOS, or a source installation, use the installation guide.

2. Download and run

The 528.5 MB portable checkpoint contains both generators. No separate stage-1 checkpoint, discriminator, VGG, or LPIPS weights are needed for inference.

mkdir -p checkpoints examples
curl -fL https://huggingface.co/KumaPower/CSA/resolve/main/csa-stage2-step97397.pt \
  -o checkpoints/csa-stage2-step97397.pt
for name in masked mask expected; do
  curl -fL "https://huggingface.co/KumaPower/CSA/resolve/main/examples/${name}.png" \
    -o "examples/${name}.png"
done

csa infer \
  --checkpoint checkpoints/csa-stage2-step97397.pt \
  --image examples/masked.png \
  --mask examples/mask.png \
  --output outputs/restored.png \
  --device auto

3. Supply your image and mask

Input image	Mask	CSA output

White = fill; black = keep. Image and mask dimensions must match. An intact image also works: the CLI masks the hole internally. The example above is a real checkpoint output.

The model was trained at 256 × 256. Other sizes are padded and cropped, but their restoration quality has not been validated. --device auto selects CUDA, then MPS, then CPU. See usage for custom inputs and folder evaluation.

Train and extend

Use the source repository and two-stage training guide to prepare image indices, train rough prediction, initialize refinement, and resume a run. Install the training dependencies with python -m pip install -e '.[perceptual]' from the checkout.

Objective	Fresh-training recipe
Stage 1	0.1 × L1 + 1 × LPIPS
Stage 2	0.1 × L1 + 1 × LPIPS + 0.01 × consistency + 0.002 × (image GAN + feature GAN)

Stage 2 uses original PatchGAN architectures with random initialization and RaLSGAN losses. Rectangles, strokes, or their union cover 10–40% of each image; masks are generated online. The saved configuration records the selected experiment.

Hardware and batch size

Setting	Stage 1 LPIPS phase	Stage 2 refinement
GPU / VRAM capacity	1 × NVIDIA H100 80GB HBM3	1 × NVIDIA H100 80GB HBM3
Optimizer batch	256	192
Precision	BF16 mixed precision	BF16 mixed precision
Image resolution	256 × 256	256 × 256
LPIPS activation chunk	8 images	8 images
Observed PyTorch peak allocated / reserved	9.15 / 12.08 GiB, training snapshot	31.18 / 34.42 GiB, two-update probe

80 GB is capacity; the last row is measured usage. LPIPS chunking is separate from optimizer batch size. The guide starts at batch 8 and provides a memory probe for your hardware.

Data and official sources

Dataset	Used in this release	Where to find it
Unsplash-derived photographs	296,145 training records; 2,000 held-out records	Unsplash Dataset · access and schema
Places365-Standard	901,730 images: half of 1,803,460 training images	Download page · CSAIL repository
HDTF talking-head frames	Added during stage 2: 30,128 frames represented by 299,469 sampling entries	HDTF data and preprocessing

The base mixture contains 1,197,875 entries; the final stage-2 mixture contains 1,497,344 sampling entries, approximately 20% faces. Repeated entries are not unique images. The exact local splits are not distributed; dataset preparation and provenance explain the selection and how to use public sources or your own images.

Checkpoint and provenance

The portable artifact is for inference. It contains no optimizer or discriminator state and cannot resume training. The full training checkpoint is preserved separately and is not included in this release; fresh training is documented in the guide.

The model retains width 64, depth 8, and a tanh refinement head. Export preserves all generator tensors; see export verification, software validation, and checksums. The evaluation record and selected preview preserve the selected checkpoint's evidence.

Citation and license

@InProceedings{Liu_2019_ICCV,
  author    = {Liu, Hongyu and Jiang, Bin and Xiao, Yi and Yang, Chao},
  title     = {Coherent Semantic Attention for Image Inpainting},
  booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision},
  month     = {October},
  year      = {2019}
}

The release retains the repository's CC BY-NC 4.0 / noncommercial research-use notice. Dataset and example-image rights remain with their respective owners.