| """Command-line inference for the Ermolaev historical print style model.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from huggingface_hub import hf_hub_download |
| from PIL import Image, ImageOps |
|
|
| from model import load_generator |
|
|
|
|
| REPO_ID = "tuva/Ermolaev" |
| DEFAULT_WEIGHTS = "weights/ermolaev_epoch2_recommended.pth" |
|
|
|
|
| def stylize(source, destination, weights=DEFAULT_WEIGHTS, max_side=1600, device=None): |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| weights_path = hf_hub_download(REPO_ID, weights) |
| generator = load_generator(weights_path, device) |
|
|
| with Image.open(source) as raw: |
| image = ImageOps.exif_transpose(raw).convert("RGB") |
| scale = min(1.0, max_side / max(image.size)) |
| if scale < 1: |
| image = image.resize(tuple(round(v * scale) for v in image.size), Image.Resampling.LANCZOS) |
| array = np.asarray(image, dtype=np.float32).copy() / 127.5 - 1.0 |
| tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).to(device) |
| height, width = tensor.shape[-2:] |
| pad_h, pad_w = (-height) % 4, (-width) % 4 |
| if pad_h or pad_w: |
| tensor = F.pad(tensor, (0, pad_w, 0, pad_h), mode="reflect") |
| with torch.inference_mode(): |
| output = generator(tensor)[0, :, :height, :width] |
| output = ((output.clamp(-1, 1) + 1) * 127.5).permute(1, 2, 0).byte().cpu().numpy() |
| result = Image.fromarray(output) |
| destination = Path(destination) |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| result.save(destination) |
| return destination |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("input") |
| parser.add_argument("output") |
| parser.add_argument("--weights", default=DEFAULT_WEIGHTS) |
| parser.add_argument("--max-side", type=int, default=1600) |
| args = parser.parse_args() |
| print(stylize(args.input, args.output, args.weights, args.max_side)) |
|
|
