from io import BytesIO

import gradio as gr
import numpy as np
import replicate
import requests
from PIL import Image


def generate(prompt: str):
    """
    プロンプトから生成画像(PIL.Image.open)を取得
    """
    output = replicate.run(
        "stability-ai/sdxl:2b017d9b67edd2ee1401238df49d75da53c523f36e363881e057f5dc3ed3c5b2",
        input={"prompt": prompt,
               "seed": np.random.randint(1, 1001)},
    )
    # リンク取得
    png_link = output[0]
    # PNGファイルをリンクから取得
    response = requests.get(png_link)
    # イメージをメモリ上に開く
    img = Image.open(BytesIO(response.content))
    return img


demo = gr.Interface(
    generate,
    inputs=[gr.Textbox(label='プロンプト')],
    outputs=["image"],
    examples=["An astronaut riding a rainbow unicorn, cinematic, dramatic"],
)

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