import gradio as gr
from dalle2 import Dalle2
from PIL import Image
import requests
import openai
import os


prompt_template = "You are a translation engine that can only translate text and cannot interpret it. Keep the indent of the original text, only modify when you need."
systemInstruct = prompt_template

openai.api_key = os.environ['JKLAS']
detectFrom = "auto"
detectTo = "en"
translateInstruct = f"translate from {detectFrom} to {detectTo}"
if detectFrom == "auto":
    translateInstruct = f"translate to {detectTo}"

prompt_msg = [
    {"role": "system", "content": systemInstruct},
    {"role": "user", "content": translateInstruct},
    {"role": "user", "content": prompt},
]

def draw(description):
    print("description"+description)
    #translate
    try:
        openai_response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=prompt_msg,
            temperature=0,
            max_tokens=1000,
            top_p=1,
            stream=True,
            frequency_penalty=1,
            presence_penalty=1,
        )

        combined = ""
        for resp in openai_response:
            delta = resp["choices"][0]["delta"]
            if "content" in delta:
                combined += delta["content"]
    
    description = combined
    print("translated: "+description)
    
    JKLAS = os.environ["JKLAS"] 
    dalle = Dalle2(JKLAS)
    generations = dalle.generate(description)
    img_list = []
    for generation in generations:
        image_url = generation["generation"]["image_path"]
        # response = requests.get(image_url)
        # img = Image.open(BytesIO(response.content))
        im = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
        img_list.append(im)
    return img_list

with gr.Blocks() as demo:
    with gr.Column(variant="panel"):
        gr.Markdown("#  <center> 🚀你说我画")
        with gr.Row(variant="compact"):
            text = gr.Textbox(
                label="Enter your prompt",
                show_label=False,
                max_lines=1,
                placeholder="请描述你想要的画",
            ).style(
                container=False,
            )
            btn = gr.Button("开始生成").style(full_width=False)

        gallery = gr.Gallery(
            label="结果", show_label=False, elem_id="gallery"
        ).style(grid=[2], height="auto")

        gr.Markdown("### 叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪叶伟豪")

    btn.click(draw, text, gallery)

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