import requests
import gradio as gr

# 修改 greet 函数以处理图像生成
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:
        return response.content
    else:
        raise Exception(str(response.json()))

# 构建 Gradio 界面
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            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")
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Describe what you do not wish to see.")
            seed = gr.Number(label="Seed", value=0, minimum=0, maximum=4294967294, step=1)
            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(server_port=7588)