import os
from io import BytesIO

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


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


examples = [
    # ["An astronaut riding a rainbow unicorn, cinematic, dramatic", ""],
    # ["A robot painted as graffiti on a brick wall. a sidewalk is in front of the wall, and grass is growing out of cracks in the concrete.", ""],
    # ["Panda mad scientist mixing sparkling chemicals, artstation.", ""],
    ["station", ""],
    ["station, ghibli style", ""],
    ["Elon Musk", ""],
    ["Elon Musk playing Shogi", ""],
    ]


demo = gr.Interface(
    generate,
    inputs=[gr.Textbox(label='プロンプト'), gr.Textbox(label='パスワード')],
    outputs=["image"],
    examples=examples,
    allow_flagging=None,
    title="Stable Diffusion XL (SDXL 1.0)"
)

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