import os
import gradio as gr
from openai import OpenAI

# 读取 API 密钥和端点
client = OpenAI(
    api_key=os.getenv("stepfun_key"),
    base_url="https://platform.stepfun.com/v1"
)

def chat_with_step3(image, question):
    """上传图片和文本，通过 API 请求 step3 推理。"""
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": image}},
                {"type": "text", "text": question},
            ],
        }
    ]
    completion = client.chat.completions.create(
        model="step3", messages=messages, max_tokens=4096
    )
    return completion.choices[0].message.content

iface = gr.Interface(
    fn=chat_with_step3,
    inputs=[gr.Image(type="pil", label="Upload image"), gr.Textbox(label="Question")],
    outputs="text",
    title="Step3 FP8 (API) Demo",
    description="使用 StepFun 提供的 OpenAI 兼容 API 调用 step3‑fp8 模型。"
)

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