import gradio as gr
from diffusers import DiffusionPipeline

# 加载模型
pipeline = DiffusionPipeline.from_pretrained("shuttleai/shuttle-3-diffusion-fp8")

# 定义生成函数
def generate_image(prompt):
    image = pipeline(prompt).images[0]
    return image

# 创建 Gradio 界面
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="输入提示词"),
    outputs=gr.Image(label="生成的图像"),
    title="Shuttle-3 Diffusion FP8 图像生成",
    description="输入提示词，生成图像。"
)

# 启动应用
interface.launch()