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

# 修改 generate_image 函数以处理图像生成并保存为临时文件
def generate_image(api_key, prompt, aspect_ratio, model, negative_prompt, seed):
    response = requests.post(
        "https://api.stability.ai/v2beta/stable-image/generate/sd3",
        headers={
            "authorization": f"Bearer {api_key}",
            "accept": "image/*"
        },
        files={"none": ''},
        data={
            "prompt": prompt,
            "aspect_ratio": aspect_ratio,
            "model": model,
            "negative_prompt": negative_prompt,
            "seed": seed,
            "output_format": "jpeg",
        },
    )
    
    if response.status_code == 200:
        # 将 bytes 数据转换为 PIL Image 对象
        image_bytes = io.BytesIO(response.content)
        image = Image.open(image_bytes)
        # 保存为临时文件
        temp_path = "./temp_image.jpeg"
        image.save(temp_path)
        return temp_path
    else:
        raise Exception(str(response.json()))

# 构建 Gradio 界面
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            api_key_text = gr.HTML("<a href='https://platform.stability.ai/account/keys' target='_blank'>https://platform.stability.ai/account/keys</a>")
            api_key = gr.Textbox(label="API Key")
            prompt = gr.Textbox(label="Prompt", placeholder="What you wish to see in the output image.")
            aspect_ratio = gr.Dropdown(label="Aspect Ratio", choices=["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"], value="1:1")
            model = gr.Dropdown(label="Model", choices=["sd3", "sd3-turbo"], value="sd3")
            seed = gr.Number(label="Seed (Enter 0 for random seed)", value=0, minimum=0, maximum=4294967294, step=1)
            seed.description = "输入0则随机种子"
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Describe what you do not wish to see.")
            generate_btn = gr.Button("Generate Image")
        with gr.Column(scale=1):
            output = gr.Image()



            
    generate_btn.click(fn=generate_image, 
                        inputs=[api_key, prompt, aspect_ratio, model, negative_prompt, seed], 
                        outputs=output)

demo.launch()