215 lines
7.7 KiB
Python
215 lines
7.7 KiB
Python
import os
|
||
import base64
|
||
import io
|
||
|
||
from dotenv import load_dotenv
|
||
import requests
|
||
import gradio as gr
|
||
from PIL import Image
|
||
|
||
load_dotenv()
|
||
|
||
API_URL = os.getenv("IMAGE_API_URL", "https://llm.dev.maimaiag.com/v1/images/generations")
|
||
API_KEY = os.getenv("IMAGE_API_KEY", "")
|
||
|
||
|
||
def _call_remote_api(prompt: str, size: str, seed: int, num_inference_steps: int) -> Image.Image:
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
}
|
||
if API_KEY:
|
||
headers["Authorization"] = f"Bearer {API_KEY}"
|
||
|
||
payload = {
|
||
"model": "openai/z-image-turbo",
|
||
"prompt": prompt,
|
||
"size": size,
|
||
"seed": seed,
|
||
"num_inference_steps": num_inference_steps,
|
||
}
|
||
|
||
response = requests.post(API_URL, headers=headers, json=payload, timeout=300)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
|
||
b64_image = data["data"][0]["b64_json"]
|
||
image_bytes = base64.b64decode(b64_image)
|
||
return Image.open(io.BytesIO(image_bytes))
|
||
|
||
|
||
def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
|
||
"""Generate an image from the given prompt via remote API."""
|
||
if randomize_seed:
|
||
import random
|
||
seed = random.randint(0, 2**32 - 1)
|
||
|
||
size = f"{int(width)}x{int(height)}"
|
||
image = _call_remote_api(
|
||
prompt=prompt,
|
||
size=size,
|
||
seed=int(seed),
|
||
num_inference_steps=int(num_inference_steps),
|
||
)
|
||
return image, seed
|
||
|
||
# 示例提示词
|
||
examples = [
|
||
["一张黑白照片（约1950年代），捕捉了一场欢乐的跨代家庭感恩节或节日晚餐。中心人物是一位身穿白衬衫、系深色领带的微笑男子，他正在餐桌主位上兴致勃勃地切割一只大烤火鸡。他被一大群各年龄段的家庭成员包围，所有人都迫不及待地伸出盘子来接食物。许多不同年龄的孩子围在周围，眼睛睁得大大的，充满期待，有的站着，有的坐着。还有几位女性和青少年，有的在帮忙上菜，有的抱着婴儿。整个场景充满了自然的互动、笑声和热闹温馨的氛围。餐桌上摆满了传统的节日菜肴。背景是一面简单的、可能贴着壁纸的墙壁，反映了真实家庭住宅的质朴感。"],
|
||
["一幅全景场景，背景从石器时代过渡到未来（洞穴→金字塔→城堡→工厂→摩天大楼→漂浮城市），主体是前景中同一张面孔/同一个人，从左到右佩戴着符合时代特征的头盔：骨/兽皮头饰、青铜古代头盔、中世纪板甲头盔、一战钢盔、现代太空头盔，以及未来能量/全息头盔。"],
|
||
["科幻电影的电影海报，宇航员站在火星上凝望地球，戏剧性的灯光，黑暗的氛围，顶部有大号金属质感字体标题\"THE VOYAGE\""],
|
||
["兴奋的年轻内容创作者指着漂浮的机器人全息图，充满活力的紫色和蓝色霓虹背景，高对比度，YouTube缩略图风格，粗体黄色文字\"AI UPDATE\""],
|
||
["现代工作空间，木质桌面上放着笔记本电脑和咖啡，晨光透过窗户洒入，温馨的氛围，照片级真实感，高质量，便利贴上写着\"Productivity\""],
|
||
]
|
||
|
||
# 自定义主题，深色科技感 (Gradio 6)
|
||
custom_theme = gr.themes.Soft(
|
||
primary_hue="cyan",
|
||
secondary_hue="violet",
|
||
neutral_hue="zinc",
|
||
font=gr.themes.GoogleFont("Inter"),
|
||
text_size="lg",
|
||
spacing_size="md",
|
||
radius_size="lg"
|
||
).set(
|
||
button_primary_background_fill="*primary_500",
|
||
button_primary_background_fill_hover="*primary_400",
|
||
button_primary_text_color="*neutral_950",
|
||
block_title_text_weight="600",
|
||
block_background_fill="*neutral_900",
|
||
body_background_fill="*neutral_950",
|
||
body_text_color="*neutral_200",
|
||
block_label_text_color="*neutral_300",
|
||
input_background_fill="*neutral_800",
|
||
border_color_primary="*neutral_700",
|
||
)
|
||
|
||
# 构建 Gradio 界面
|
||
with gr.Blocks(fill_height=True) as demo:
|
||
# 页眉
|
||
gr.Markdown(
|
||
"""
|
||
# Z-Image-Turbo
|
||
极速 AI 图像生成 • 仅需 8 步即可生成精美图像
|
||
""",
|
||
elem_classes="header-text"
|
||
)
|
||
|
||
with gr.Row(equal_height=False):
|
||
# 左栏 - 输入控制
|
||
with gr.Column(scale=1, min_width=320):
|
||
prompt = gr.Textbox(
|
||
label="提示词",
|
||
placeholder="描述你想要生成的图像...",
|
||
lines=5,
|
||
max_lines=10,
|
||
autofocus=True,
|
||
)
|
||
|
||
with gr.Accordion("高级设置", open=False):
|
||
with gr.Row():
|
||
height = gr.Slider(
|
||
minimum=512,
|
||
maximum=2048,
|
||
value=1024,
|
||
step=64,
|
||
label="高度",
|
||
info="图像高度（像素）"
|
||
)
|
||
width = gr.Slider(
|
||
minimum=512,
|
||
maximum=2048,
|
||
value=1024,
|
||
step=64,
|
||
label="宽度",
|
||
info="图像宽度（像素）"
|
||
)
|
||
|
||
num_inference_steps = gr.Slider(
|
||
minimum=1,
|
||
maximum=20,
|
||
value=9,
|
||
step=1,
|
||
label="推理步数",
|
||
info="9 步 = 8 次 DiT 前向传播（推荐）"
|
||
)
|
||
|
||
with gr.Row():
|
||
randomize_seed = gr.Checkbox(
|
||
label="随机种子",
|
||
value=True,
|
||
)
|
||
seed = gr.Number(
|
||
label="种子",
|
||
value=42,
|
||
precision=0,
|
||
visible=False,
|
||
)
|
||
|
||
def toggle_seed(randomize):
|
||
return gr.Number(visible=not randomize)
|
||
|
||
randomize_seed.change(
|
||
toggle_seed,
|
||
inputs=[randomize_seed],
|
||
outputs=[seed]
|
||
)
|
||
|
||
generate_btn = gr.Button(
|
||
"生成图像",
|
||
variant="primary",
|
||
size="lg",
|
||
scale=1
|
||
)
|
||
|
||
# 示例提示词
|
||
gr.Examples(
|
||
examples=examples,
|
||
inputs=[prompt],
|
||
label="试试这些提示词",
|
||
examples_per_page=5,
|
||
)
|
||
|
||
# 右栏 - 输出
|
||
with gr.Column(scale=1, min_width=320):
|
||
output_image = gr.Image(
|
||
label="生成的图像",
|
||
type="pil",
|
||
format="png",
|
||
show_label=False,
|
||
height=600,
|
||
buttons=["download"],
|
||
)
|
||
|
||
used_seed = gr.Number(
|
||
label="使用的种子",
|
||
interactive=False,
|
||
container=True,
|
||
)
|
||
|
||
# 连接生成按钮
|
||
generate_btn.click(
|
||
fn=generate_image,
|
||
inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
|
||
outputs=[output_image, used_seed],
|
||
)
|
||
|
||
# 也支持在提示词框中按回车键生成
|
||
prompt.submit(
|
||
fn=generate_image,
|
||
inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
|
||
outputs=[output_image, used_seed],
|
||
)
|
||
|
||
if __name__ == "__main__":
|
||
demo.launch(
|
||
server_name="0.0.0.0",
|
||
server_port=7860,
|
||
theme=gr.themes.Soft(primary_hue="orange"),
|
||
css="""
|
||
|
||
""",
|
||
footer_links=[
|
||
],
|
||
mcp_server=False
|
||
)
|