import gradio as gr
from transformers import pipeline

pipe_aesthetic = pipeline("image-classification", "cafeai/cafe_aesthetic")
def aesthetic(input_img):
    data = pipe_aesthetic(input_img, top_k=2)
    final = {}
    for d in data:
        final[d["label"]] = d["score"]
    return final
demo_aesthetic = gr.Interface(fn=aesthetic, inputs=gr.Image(type="pil"), outputs=gr.Label(label="aesthetic"))

pipe_style = pipeline("image-classification", "cafeai/cafe_style")
def style(input_img):
    data = pipe_style(input_img, top_k=5)
    final = {}
    for d in data:
        final[d["label"]] = d["score"]
    return final
demo_style = gr.Interface(fn=style, inputs=gr.Image(type="pil"), outputs=gr.Label(label="style"))

pipe_waifu = pipeline("image-classification", "cafeai/cafe_waifu")
def waifu(input_img):
    data = pipe_waifu(input_img, top_k=5)
    final = {}
    for d in data:
        final[d["label"]] = d["score"]
    return final
demo_waifu = gr.Interface(fn=waifu, inputs=gr.Image(type="pil"), outputs=gr.Label(label="waifu"))

gr.Parallel(demo_aesthetic, demo_style, demo_waifu).launch()