TareHimself/AnimeMangaInpainting-torchscript

AnimeMangaInpainting - TorchScript

A TorchScript build of dreMaz/AnimeMangaInpainting (lama_large_512px.ckpt) - "big-lama" fine-tuned on 300k manga/anime images - so consumers can inpaint without carrying any model code.

import torch
from huggingface_hub import hf_hub_download

model = torch.jit.load(hf_hub_download("TareHimself/AnimeMangaInpainting-torchscript",
                                       "anime_manga_lama.pt")).eval()

# image: (N,3,H,W) float32 RGB in [0,1], H&W multiple of 8
# mask:  (N,1,H,W) float32 in [0,1]; >= 0.5 marks pixels to erase
out = model(image, mask)          # (N,3,H,W) float32 in [0,1]

Files

file	what
anime_manga_lama.pt	the TorchScript module (traced + frozen). This is all a consumer needs.
lama_large_512px.ckpt	the original checkpoint, unchanged, kept as the source of truth
to_torchscript.py	reproducible .ckpt -> .pt conversion + numerical check against eager
config.json	Hub-readable architecture/contract sidecar (mirrors the module's embedded meta.json)
lama_ffc.py	the FFC generator class (trimmed from advimman/lama), for the raw-weights path
inference.py	pad_to_modulo helper + single-image and batch examples
pyproject.toml / uv.lock	pinned environment for reproducing the export (uv)

forward(image, mask) contract

arg	shape	dtype	notes
image	(N, 3, H, W)	float32	RGB, values in [0, 1], H and W must be multiples of 8
mask	(N, 1, H, W)	float32	values in [0, 1]; >= 0.5 marks pixels to erase
return	(N, 3, H, W)	float32	values in [0, 1]

Baked into the graph: mask binarisation, masked-image construction, the 4-channel assembly LaMa expects, an output clamp, and the mask * predicted + (1 - mask) * image composite (pixels outside the mask are returned untouched).

Not baked in, on purpose:

• Padding to a multiple of 8. Do it before the call and crop the result back - inference.py:pad_to_modulo does this. Left to the caller because it is where peak VRAM and FFT plan cost are decided.
• Colour conversion. Pass RGB. OpenCV images are BGR - convert first.
• uint8 normalisation. Pass float in [0, 1].

Batching

Stack same-size inputs on dim 0. For varying sizes, pad each to a common multiple-of-8 size and crop the results back (inference.py:inpaint_batch). Peak VRAM is N x maxH x maxW at the Fourier layers, so bucket by size before batching. LaMa's FFT is also fastest at sizes with small prime factors and rebuilds a cuFFT plan per new (H, W); for mixed workloads prefer a few fixed sizes (512 / 768 / 1024) over "next multiple of 8".

Raw-weights path (no TorchScript)

import torch
from lama_ffc import load_generator

gen = load_generator("lama_large_512px.ckpt")     # strict load, eval mode
# gen expects the 4-channel [0,1] tensor cat([image*(1-mask), mask], dim=1)
# and returns the predicted image in [0,1] with no compositing.

This keeps torch.compile / torch.export open for anyone who wants them.

Reproduce

uv sync
uv run python to_torchscript.py --ckpt lama_large_512px.ckpt --out anime_manga_lama.pt

The export traces the model and asserts the traced output matches eager to < 1e-3 across several resolutions and batch sizes. Embedded metadata:

extra = {"meta.json": ""}
torch.jit.load("anime_manga_lama.pt", _extra_files=extra)
print(extra["meta.json"])

Versions

Exported with torch 2.11.0 (pinned in uv.lock). torch >= 2.1 is expected to load it; if torch.jit.load fails on an older build, re-run to_torchscript.py with your torch.

Attribution and licence

• Checkpoint: dreMaz/AnimeMangaInpainting - MIT
• Architecture (lama_ffc.py): advimman/lama - Apache-2.0 (see NOTICE)
• LaMa: Suvorov et al., Resolution-robust Large Mask Inpainting with Fourier Convolutions, WACV 2022 - https://arxiv.org/abs/2109.07161

This repo is MIT-licensed; the vendored architecture code remains under Apache-2.0 as noted in NOTICE.