'''Image Completion Demo (ImageGPT)

- 2021-12-10
    - taken from https://huggingface.co/spaces/nielsr/imagegpt-completion
'''

from PIL import Image
import matplotlib.pyplot as plt

import os
import requests
from glob import glob
import gradio as gr
from loguru import logger

import torch
from transformers import ImageGPTFeatureExtractor, ImageGPTForCausalImageModeling

# ========== Settings ==========
EXAMPLE_DIR = 'examples'
examples = sorted(glob(os.path.join(EXAMPLE_DIR, '*.jpg')))

# ========== Logger ==========
logger.add('app.log', mode='a')
logger.info('===== APP RESTARTED =====')

# ========== Models ==========
feature_extractor = ImageGPTFeatureExtractor.from_pretrained(
    "openai/imagegpt-medium")
model = ImageGPTForCausalImageModeling.from_pretrained(
    "openai/imagegpt-medium")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)


def process_image(image):
    logger.info('--- image file received')
    return image.name


iface = gr.Interface(
    process_image,
    title="이미지의 절반을 지우고 절반을 채워 넣어주는 Image Completion 데모입니다 (ImageGPT)",
    description='주어진 이미지의 절반 아래를 AI가 채워 넣어줍니다',
    inputs=gr.inputs.Image(type="pil"),
    outputs=gr.outputs.Image(type="pil", label="Model input + completions"),
    examples=examples,
    enable_queue=True,
    article='<p style="text-align:center">i-Scream AI</p>',
)
iface.launch()