import json
import gradio as gr

from coze import call_plugin


async def generate_image(query):
    parameters = {
        "text_prompts": [{"text": query}],
        # Stable Diffusion XL v1.0
        "model_type": 2,
        "width": 1344,
        "height": 768,
    }
    async for message in call_plugin(parameters, "Stable Diffusion:textToImage"):
        event = message.get("event")
        data = message.get("data")

        if event == "conversation.message.completed" and data:
            try:
                data = json.loads(data)  # 尝试解析 JSON 数据
                response_type = data.get("type")
                if response_type == "tool_response":
                    content = json.loads(data.get("content"))
                    image_url = content.get("data").get("images")[0]
                    if image_url:
                        yield f'<img src="{image_url})" />'
            except json.JSONDecodeError:
                yield "Error: Failed to decode JSON data."


with gr.Blocks() as demo:
    prompt = gr.Textbox(label="输入描述")
    output = gr.HTML()
    generate_button = gr.Button("生成图片")
    generate_button.click(fn=generate_image, inputs=prompt, outputs=output)

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