app.py

import gradio as gr
from setup import setup
import torch
import gc
from PIL import Image
from manga_line_extraction.model import MangaLineExtractor
from anime2sketch.model import Anime2Sketch

setup()

print("Setup finished")


def flush():
    gc.collect()
    torch.cuda.empty_cache()


def extract(image):
    extractor = MangaLineExtractor("./models/erika.pth", "cpu")
    result = extractor.predict(image)
    del extractor
    flush()
    return result


def convert_to_sketch(image):
    to_sketch = Anime2Sketch("./models/netG.pth", "cpu")
    result = to_sketch.predict(image)
    del to_sketch
    flush()
    return result


def start(image):
    return [extract(image), convert_to_sketch(Image.fromarray(image).convert("RGB"))]


def clear():
    return [None, None]


def ui():
    with gr.Blocks() as blocks:
        gr.Markdown(
            """
        # Anime to Sketch 
        Unofficial demo for converting illustrations into sketches. 
        Original repos:
        - [MangaLineExtraction_PyTorch](https://github.com/ljsabc/MangaLineExtraction_PyTorch)
        - [Anime2Sketch](https://github.com/Mukosame/Anime2Sketch)
        """
        )

        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="Input", interactive=True)

                extract_btn = gr.Button("Start", variant="primary")
                clear_btn = gr.Button("Clear", variant="secondary")

            with gr.Column():
                # with gr.Row():
                extract_output_img = gr.Image(
                    label="MangaLineExtraction", interactive=False
                )
                to_sketch_output_img = gr.Image(label="Anime2Sketch", interactive=False)

        gr.Examples(
            fn=start,
            examples=[
                ["./examples/1.jpg"],
                ["./examples/2.jpg"],
                ["./examples/3.jpg"],
                ["./examples/4.jpg"],
                ["./examples/5.jpg"],
                ["./examples/6.jpg"],
                ["./examples/7.jpg"],
                ["./examples/8.jpg"],
            ],
            inputs=[input_img],
            outputs=[extract_output_img, to_sketch_output_img],
            label="Examples",
            # cache_examples=True,
        )

        gr.Markdown("Images are from nijijourney.")

        extract_btn.click(
            fn=start,
            inputs=[input_img],
            outputs=[extract_output_img, to_sketch_output_img],
        )

        clear_btn.click(
            fn=clear,
            inputs=[],
            outputs=[extract_output_img, to_sketch_output_img],
        )

    return blocks


if __name__ == "__main__":
    ui().launch()