import os
import io
import base64
import gradio as gr
from PIL import Image
from openai import OpenAI

def pil_to_data_url(img: Image.Image, fmt="PNG"):
    buf = io.BytesIO()
    img.save(buf, format=fmt)
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    mime = "image/png" if fmt.upper() == "PNG" else "image/jpeg"
    return f"data:{mime};base64,{b64}"

client = OpenAI(
    api_key=os.getenv("STEPFUN_KEY"),  # 在 HF Secrets 里配置
    base_url="https://platform.stepfun.com/v1",
)

def chat_with_step3(image: Image.Image, question: str):
    if image is None:
        return "请先上传图片。"
    if not question:
        question = "请描述这张图片。"

    data_url = pil_to_data_url(image, fmt="PNG")

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": data_url}},
                {"type": "text", "text": question},
            ],
        }
    ]
    resp = client.chat.completions.create(
        model="step3-fp8",  # 更稳妥的模型名
        messages=messages,
        max_tokens=1024,
    )
    return resp.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__":
    # HF Spaces 内不必 share=True；保持默认即可
    iface.launch()