app.py

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

API_URL = os.getenv("API_URL")
API_KEY = os.getenv("API_KEY")

SIZE_OPTIONS = ["1024x1024", "848x1264", "1264x848", "1376x768", "1200x896", "896x1200", "768x1376"]


def generate_image(prompt: str, size: str, seed: int, use_pe: bool):
    if not prompt.strip():
        raise gr.Error("请输入图像描述提示词")

    headers = {
        "Authorization": f"bearer {API_KEY}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": "ernie-image-turbo",
        "prompt": prompt,
        "n": int(1),
        "response_format": "b64_json",
        "size": size,
        "seed": int(seed),
        "use_pe": use_pe,
        "num_inference_steps": 8,
        "guidance_scale": 1.0
    }

    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=300)
    except requests.exceptions.Timeout:
        raise gr.Error("请求超时，请稍后重试")
    except requests.exceptions.RequestException as e:
        raise gr.Error(f"网络请求失败: {e}")

    # 先解析响应体，再判断状态码，以便获取 API 返回的具体错误信息
    try:
        data = response.json()
    except Exception:
        raise gr.Error(f"响应解析失败，HTTP {response.status_code}: {response.text[:200]}")

    if not response.ok:
        err = data.get("error", {})
        msg = err.get("message") or err.get("msg") or response.text[:200]
        raise gr.Error(f"API 错误 ({response.status_code}): {msg}")

    if "error" in data:
        err = data["error"]
        raise gr.Error(f"API 错误: {err.get('message', err)}")

    # 响应格式：{"created": ..., "data": [{"b64_json": "...", "revised_prompt": "...", "url": null}], "trace_id": "..."}
    items = data.get("data", [])
    if not items:
        raise gr.Error("未能获取到生成的图像，请检查 API 响应")

    images = []
    revised_prompts = []
    for item in items:
        b64 = item.get("b64_json", "")
        if b64:
            img_bytes = base64.b64decode(b64)
            images.append(Image.open(io.BytesIO(img_bytes)))
        rp = item.get("revised_prompt", "")
        if rp:
            revised_prompts.append(rp)

    if not images:
        raise gr.Error("未能解析图像数据，请检查 API 响应")

    revised_text = "\n\n".join(revised_prompts) if revised_prompts else ""
    return images[0], revised_text


with gr.Blocks(title="ERNIE-Image-Turbo 8B") as demo:
    gr.Markdown(
        """
        # 🎨 ERNIE-Image-Turbo 8B
        A demo for Baidu ERNIE-Image-Turbo image generation. Enter a description in Chinese or English to generate high-quality images with a single click.
        """,
        elem_classes="title-text",
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="例如：一只黑猫，蓝帽子，写实风格",
                lines=4,
                max_lines=8,
            )
            with gr.Row():
                size_dropdown = gr.Dropdown(
                    label="Image Size",
                    choices=SIZE_OPTIONS,
                    value="1024x1024",
                )
                # n_slider = gr.Slider(
                #     label="生成数量（1次只能1张图）",
                #     minimum=1,
                #     maximum=1,
                #     step=1,
                #     value=1,
                # )
            seed_number = gr.Number(
                label="Random Seed（-1: Random）",
                value=-1,
                precision=0,
            )
            use_pe_checkbox = gr.Checkbox(
                label="use_pe",
                value=True,
            )
            generate_btn = gr.Button(
                "Generate",
                variant="primary",
                elem_classes="generate-btn",
            )

        with gr.Column(scale=2):
            # output_gallery = gr.Gallery(
            #     label="生成结果",
            #     show_label=True,
            #     type="pil",
            #     height=300,
            #     columns=4,
            #     object_fit="cover",
            #     preview=True,
            #     selected_index=None,
            # )
            output_gallery = gr.Image(
                label="Output",
                type="pil",
                height=400,
            )
            revised_prompt_output = gr.Textbox(
                label="revised_prompt",
                lines=4,
                interactive=False,
            )

    gr.Examples(
        examples=[
            ["这是一张人像摄影图片。画面中心是一位年轻女性，正面对着镜头。她留着深棕色的长发，梳成了两条垂在双肩前的麻花辫，发尾用深色发圈扎紧，额前留着轻薄的空气刘海。她头戴一顶浅驼色的贝雷帽，身穿一件厚实的米白色粗棒针织毛衣，内搭白色高领衬衫。女子的右手抬起，手指轻触嘴唇附近，眼神直视镜头，表情显得有些惊讶或羞涩。 背景环境较为模糊，具有明显的景深效果。左侧可见一个带有金属灯座的台灯和散落的旧纸张，看起来像是一份乐谱。背景中整齐排列着几本厚重的精装书籍，书脊呈褐色，带有古典风格。右侧似乎有窗户透进来的自然光线，使得整体色调偏向温暖柔和的复古风格，光线从右侧照在人物面部，形成自然的阴影。画面整体色彩以暖棕色、米色和奶油色为主，营造出一种文艺、静谧的室内氛围。图片中未发现明显的水印、Logo或文字标记。", "848x1264", -1],
            ["雄伟的雪山，日出时分，超写实摄影", "1024x1024", 42],
            ["一张桌面端网页截图，展现了一篇关于科技博客的文章页面。顶部是浅灰色的浏览器界面，左上角有红黄绿三个圆形控制按钮。 当前激活的标签页文字为 'Understanding Diffusion Models'，下方地址栏中清晰显示着 URL：'https://www.ai-research-blog.com/understanding-diffusion-models'。 网页整体采用白底、深灰字体的现代简约设计，带有蓝色的强调元素。 网页顶部的导航栏左侧有一个蓝色的神经网络节点图标以及品牌名 'AI Research'，右侧横向排列着四个导航菜单项：'Home'、'Articles'、'Tutorials' 和 'About'。 网页主体采用左右分栏布局，左侧较宽的主内容区顶部是一个醒目的黑色粗体大标题：'深入浅出：扩散模型 (Diffusion Models) 原理剖析'。大标题下方是正文段落：'扩散模型是一类强大的生成模型，近年来在图像生成上取得了革命性的成功，例如 Stable Diffusion 和 Midjourney。其核心原理主要由两个阶段构成：前向扩散过程（逐步加噪）和反向生成过程（逐步去噪）。' 正文下方是一张居中放置、带有浅灰色边框的原理说明图表。 图表从左到右展示“前向过程”，最左侧是一张清晰的柴犬照片，下方标注 '$X_0$'；随着向右的灰色箭头引导，照片依次演变，噪点逐渐增多，中间的过渡图片下方标注 '$X_t$'，直到最右侧变成一幅完全由彩色随机像素组成的纯噪声图，下方标注 '$X_T$'。 这排图像的上方有一条弧线箭头，文字标注为 '前向过程 (Forward Process): 逐步添加高斯噪声'。 下方有一条弧线箭头，文字标注为 '反向过程 (Reverse Process): 神经网络 (如 U-Net) 学习预测并去除噪声'。 图表下方还有一段补充说明文字：'在实际训练中，模型并不直接生成图像，而是学习在每个时间步 $t$ 预测图像中加入的噪声分布。通过最小化真实噪声与模型预测噪声之间的误差，来实现高质量的去噪还原。' 在网页右侧较窄的侧边栏中，顶部有一个小标题 '相关阅读 (Related Articles)'，下方垂直排列着两条带有下划线的超链接文本：第一条是 '1. GAN 与 Diffusion 模型的深度对比'，第二条是 '2. 浅析 U-Net 架构在图像生成中的关键作用'。 整个页面排版结构清晰，信息量丰富，呈现出专业学术博客的视觉风格。", "1264x848", -1],
            ["一张高清晰度的写实照片，采用正面视角拍摄了一块通透的玻璃白板。白板表面使用褪色的绿色马克笔手绘了《海贼王》中草帽海贼团（Straw Hat Pirates）的全员合影。在横向展开的宽屏画面中，从左至右依次绘制着布鲁克（Brook）、弗兰奇（Franky）、罗宾（Robin）、索隆（Zoro）、路飞（Luffy）居于中心、娜美（Nami）、山治（Sanji）、乌索普（Usopp）、乔巴（Chopper）和甚平（Jinbe）。各个角色以生动的马克笔线条勾勒出标志性的轮廓与姿态。 画面的物理细节极度逼真，绿色马克笔的墨迹呈现出明显的水分不足与褪色感，线条中间有墨水变淡、边缘有颜料堆积的真实白板画特征，部分笔画带有轻微的断续感。玻璃白板本身具有极强的真实感，表面泛着柔和的室内环境反光和高光，透过厚重的玻璃隐约可见后方被景深虚化处理的室内背景。白板上还可以察觉到之前擦拭留下的微弱绿色墨渍残留、玻璃反光面上的细小划痕和灰尘颗粒。整体画面光线均匀，色调干净，摄影级的4K高分辨率将玻璃的光学质感、褪色墨水的纹理与手绘线条的细节展现得淋漓尽致。", "1264x848", 100],
        ],
        inputs=[prompt_input, size_dropdown, seed_number, use_pe_checkbox],
        label="示例提示词",
    )

    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, size_dropdown, seed_number, use_pe_checkbox],
        outputs=[output_gallery, revised_prompt_output],
        concurrency_limit=10,
    )

if __name__ == "__main__":
    demo.queue().launch(
        theme=gr.themes.Soft(primary_hue="blue"),
        css="""
        .title-text { text-align: center; }
        .generate-btn { min-width: 160px; }
        """,
    )