| """Convert lama_large_512px.ckpt to a self-contained TorchScript module. |
| |
| python to_torchscript.py \ |
| --ckpt lama_large_512px.ckpt \ |
| --out anime_manga_lama.pt |
| |
| The exported module has no dependency on lama_ffc.py -- consumers only need |
| `torch` of a compatible version (see requirements.txt): |
| |
| import torch |
| m = torch.jit.load("anime_manga_lama.pt").eval() |
| out = m(image, mask) # see the contract below |
| |
| Contract of the exported `forward(image, mask)`: |
| |
| image : (N, 3, H, W) float32, RGB, values in [0, 1], H and W multiples of 8 |
| mask : (N, 1, H, W) float32, values in [0, 1]; >= 0.5 marks pixels to erase |
| -> (N, 3, H, W) float32, values in [0, 1] |
| |
| Baked in: mask binarisation, masked-image construction, 4-channel assembly, |
| output clamp and the `mask * predicted + (1 - mask) * image` composite. |
| NOT baked in: padding to a multiple of 8 (do it before the call and crop the |
| result back -- see inference.py), colour-space conversion, uint8 normalisation. |
| """ |
|
|
| import argparse |
| import json |
| import time |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from lama_ffc import BIG_LAMA_KWARGS, load_generator |
|
|
| SOURCE_CKPT = "dreMaz/AnimeMangaInpainting :: lama_large_512px.ckpt" |
| UPSTREAM = "https://github.com/advimman/lama (Apache-2.0)" |
|
|
|
|
| class AnimeMangaLama(nn.Module): |
| """LaMa generator + baked pre/post processing. See module docstring.""" |
|
|
| def __init__(self, generator: nn.Module): |
| super().__init__() |
| self.generator = generator |
|
|
| def forward(self, image: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: |
| mask = (mask >= 0.5).to(image.dtype) |
| masked = image * (1.0 - mask) |
| x = torch.cat([masked, mask], dim=1) |
| predicted = self.generator(x).clamp(0.0, 1.0) |
| return predicted * mask + image * (1.0 - mask) |
|
|
|
|
| def build_module(ckpt_path: str) -> AnimeMangaLama: |
| gen = load_generator(ckpt_path) |
| module = AnimeMangaLama(gen).eval() |
| for p in module.parameters(): |
| p.requires_grad_(False) |
| return module |
|
|
|
|
| @torch.inference_mode() |
| def export(ckpt_path: str, out_path: str) -> None: |
| module = build_module(ckpt_path) |
|
|
| example = (torch.rand(1, 3, 256, 256), (torch.rand(1, 1, 256, 256) > 0.5).float()) |
| traced = torch.jit.trace(module, example, check_trace=False) |
| traced = torch.jit.freeze(traced.eval()) |
|
|
| validate(module, traced) |
|
|
| meta = { |
| "name": "anime-manga-lama-torchscript", |
| "architecture": "LaMa FFCResNetGenerator (big-lama)", |
| "generator_kwargs": _jsonable(BIG_LAMA_KWARGS), |
| "source_checkpoint": SOURCE_CKPT, |
| "upstream_architecture": UPSTREAM, |
| "license": "MIT (checkpoint) / Apache-2.0 (architecture code)", |
| "torch_version": torch.__version__, |
| "exported_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), |
| "forward": "forward(image[N,3,H,W] f32 RGB [0,1], mask[N,1,H,W] f32 [0,1]) -> [N,3,H,W] f32 [0,1]", |
| "padding": "caller pads H,W to a multiple of 8 and crops the result back", |
| } |
| torch.jit.save(traced, out_path, _extra_files={"meta.json": json.dumps(meta, indent=2)}) |
| print(f"wrote {out_path}") |
| print(json.dumps(meta, indent=2)) |
|
|
| write_config(meta, out_path) |
|
|
|
|
| def write_config(meta: dict, out_path: str, path: str = "config.json") -> None: |
| """A Hub-readable sidecar. Also the query file HF uses for download counts.""" |
| cfg = { |
| "model_type": "lama", |
| "architecture": "FFCResNetGenerator", |
| "framework": "pytorch", |
| "format": "torchscript", |
| "artifact": out_path, |
| "task": "image-inpainting", |
| "generator_kwargs": meta["generator_kwargs"], |
| "input": { |
| "image": "float32 (N,3,H,W) RGB in [0,1], H and W multiples of 8", |
| "mask": "float32 (N,1,H,W) in [0,1]; >= 0.5 marks pixels to erase", |
| }, |
| "output": "float32 (N,3,H,W) RGB in [0,1] (mask*predicted + (1-mask)*image)", |
| "baked_in": ["mask binarisation", "masked-image assembly", "4-channel concat", |
| "output clamp", "composite"], |
| "not_baked_in": ["padding to multiple of 8", "colour-space conversion", |
| "uint8 normalisation"], |
| "source_checkpoint": meta["source_checkpoint"], |
| "source_checkpoint_repo": "dreMaz/AnimeMangaInpainting", |
| "upstream_architecture": meta["upstream_architecture"], |
| "exported_with_torch": meta["torch_version"], |
| "license": "MIT (checkpoint) / Apache-2.0 (architecture code)", |
| } |
| with open(path, "w") as f: |
| f.write(json.dumps(cfg, indent=2) + "\n") |
| print(f"wrote {path}") |
|
|
|
|
| def _jsonable(d): |
| return {k: (_jsonable(v) if isinstance(v, dict) else v) for k, v in d.items()} |
|
|
|
|
| @torch.inference_mode() |
| def validate(eager: nn.Module, scripted: nn.Module) -> None: |
| torch.manual_seed(0) |
| cases = [(1, 256, 256), (1, 512, 512), (2, 512, 384), (1, 1024, 768), (4, 320, 320)] |
| worst = 0.0 |
| for n, h, w in cases: |
| img = torch.rand(n, 3, h, w) |
| mask = (torch.rand(n, 1, h, w) > 0.7).float() |
| a = eager(img, mask) |
| b = scripted(img, mask) |
| d = (a - b).abs().max().item() |
| worst = max(worst, d) |
| assert a.shape == (n, 3, h, w) |
| assert b.min() >= 0.0 and b.max() <= 1.0 |
| print(f" case N{n} {h}x{w}: max|eager-scripted| = {d:.2e}") |
| assert worst < 1e-3, f"scripted output diverges from eager (max {worst:.2e})" |
| print(f" OK (worst {worst:.2e})") |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", default="lama_large_512px.ckpt") |
| ap.add_argument("--out", default="anime_manga_lama.pt") |
| args = ap.parse_args() |
| export(args.ckpt, args.out) |
|
|
