import io
from os import getenv
import gradio as gr
import requests
from PIL import Image

API_URL = (
    "https://api-inference.huggingface.co/models/ztjona/scopic-diffusion-OW-v1.4.1"
)
API_TOKEN = getenv("API_TOKEN")
headers = {"Authorization": f"Bearer {API_TOKEN}"}


def infer(prompt):
    payload = {"inputs": prompt}
    response = requests.post(API_URL, headers=headers, json=payload)
    image_bytes = response.content

    return Image.open(io.BytesIO(image_bytes))


demo = gr.Interface(
    fn=infer,
    inputs="text",
    outputs="image",
    examples=[
        ["city and clouds"],
        ["tea party"],
        ["chess player"],
        ["buddhist monk"],
        ["the man of the mask"],
    ],
    title="Generate images based on artist Oswaldo Guayasamín",
    description="Write a prompt and an image will be generated. ",
    article="""Fune tuned Stable Diffusion model. \n
    Instructions are located (here)[https://github.com/ztjona/scopic-diffusion]\n
    Check the training code in colab (here)[https://drive.google.com/file/d/1r1z8Ckqq4W9U0Wg0PcWGB_ri4DavpkCQ/view?usp=sharing]""",
    theme=gr.themes.Soft(),
)

demo.launch()