64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
import torch
|
|
|
|
def generator(pretrained=True, device="cpu", progress=True, check_hash=True):
|
|
from model import Generator
|
|
|
|
release_url = "https://github.com/bryandlee/animegan2-pytorch/raw/main/weights"
|
|
known = {
|
|
name: f"{release_url}/{name}.pt"
|
|
for name in [
|
|
'celeba_distill', 'face_paint_512_v1', 'face_paint_512_v2', 'paprika'
|
|
]
|
|
}
|
|
|
|
device = torch.device(device)
|
|
model = Generator().to(device)
|
|
|
|
if type(pretrained) == str:
|
|
# Look if a known name is passed, otherwise assume it's a URL
|
|
ckpt_url = known.get(pretrained, pretrained)
|
|
pretrained = True
|
|
else:
|
|
ckpt_url = known.get('face_paint_512_v2')
|
|
|
|
if pretrained is True:
|
|
state_dict = torch.hub.load_state_dict_from_url(
|
|
ckpt_url,
|
|
map_location=device,
|
|
progress=progress,
|
|
check_hash=check_hash,
|
|
)
|
|
model.load_state_dict(state_dict)
|
|
|
|
return model
|
|
|
|
|
|
def face2paint(device="cpu", size=512):
|
|
from PIL import Image
|
|
from torchvision.transforms.functional import to_tensor, to_pil_image
|
|
|
|
def face2paint(
|
|
model: torch.nn.Module,
|
|
img: Image.Image,
|
|
size: int = size,
|
|
side_by_side: bool = True,
|
|
device: str = device,
|
|
) -> Image.Image:
|
|
w, h = img.size
|
|
s = min(w, h)
|
|
img = img.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))
|
|
img = img.resize((size, size), Image.LANCZOS)
|
|
|
|
with torch.no_grad():
|
|
input = to_tensor(img).unsqueeze(0) * 2 - 1
|
|
output = model(input.to(device)).cpu()[0]
|
|
|
|
if side_by_side:
|
|
output = torch.cat([input[0], output], dim=2)
|
|
|
|
output = (output * 0.5 + 0.5).clip(0, 1)
|
|
|
|
return to_pil_image(output)
|
|
|
|
return face2paint
|