app.py

import argparse
import functools
import pathlib
import os
from typing import Dict, List, Optional, Tuple
import gradio as gr
import PIL.Image
from encoder import Encoder

from face_detector import FaceAligner
from generator import Generator
from huggingface_hub import hf_hub_download


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument('--models_repo_id', type=str,
                        default='senior-sigan/nijigenka')
    return parser.parse_args()


def load_examples():
    image_dir = pathlib.Path('examples')
    images = sorted(image_dir.glob('*.jpg'))
    return [[path.as_posix(), 'art'] for path in images]


def join_image_h(im1: PIL.Image.Image, im2: PIL.Image.Image) -> PIL.Image.Image:
    im1 = im1.resize(im2.size)
    dst = PIL.Image.new('RGB', (im1.width + im2.width, im1.height))
    dst.paste(im1, (0, 0))
    dst.paste(im2, (im1.width, 0))
    return dst


def predict(
    image: PIL.Image.Image,
    style: str,
    *,
    face_aligner: FaceAligner,
    encoder: Encoder,
    generator: Dict[str, Generator],
) -> Tuple[List[PIL.Image.Image], Optional[str]]:
    images = face_aligner.align(image)
    if len(images) == 0:
        error_msg = "Cannot find any face in photo"
        # gradio doesn't support empty list for images carusel, so we create dummy img
        return [PIL.Image.new('RGB', (1, 1))], error_msg

    results = []
    for img in images:
        x = encoder.predict(img)
        gen_img = generator[style].predict(x)
        result = join_image_h(img, gen_img)
        results.append(result)

    return results, None


def get_model_path(repo_id: str, filename: str):
    maybe_path = os.path.join(repo_id, filename)
    if os.path.exists(maybe_path):
        print('Using local models')
        return os.path.abspath(maybe_path)
    else:
        return hf_hub_download(
            repo_id,
            filename,
        )


def load_models(repo_id: str):
    encoder_path = get_model_path(
        repo_id,
        'encoder.onnx',
    )
    generator_art_path = get_model_path(
        repo_id,
        'face2art.onnx',
    )
    generator_anime_path = get_model_path(
        repo_id,
        'face2kuvshinov2.onnx',
    )
    shape_predictor_path = get_model_path(
        repo_id,
        'shape_predictor_68_face_landmarks.bin',
    )

    face_aligner = FaceAligner(
        image_size=512,
        shape_predictor_path=shape_predictor_path,
    )
    encoder = Encoder(model_path=encoder_path)
    generator_art = Generator(model_path=generator_art_path)
    generator_anime = Generator(model_path=generator_anime_path)

    return face_aligner, encoder, {'art': generator_art, 'anime': generator_anime}


def main():
    args = parse_args()
    gr.close_all()

    face_aligner, encoder, generator = load_models(args.models_repo_id)
    generator_types = list(generator.keys())

    func = functools.partial(
        predict,
        face_aligner=face_aligner,
        encoder=encoder,
        generator=generator,
    )
    func = functools.update_wrapper(func, predict)

    iface = gr.Interface(
        fn=func,
        inputs=[
            gr.Image(
                type='pil',
                label='Real photo with a face',
            ),
            gr.Radio(
                choices=generator_types,
                type='value',
                value=generator_types[0],
                label='Style',
            ),
        ],
        outputs=[
            gr.Gallery(label='Result'),
            gr.Textbox(label='Error'),
        ],
        examples=load_examples(),
        title='Nijigenka: Portrait to Art',
        allow_flagging='never',
    )
    iface.queue().launch()


if __name__ == '__main__':
    main()