import base64
from io import BytesIO
import os
from mistralai import Mistral
import re
from PIL import Image


api_key = os.getenv("MISTRAL_API_KEY")
Mistralclient = Mistral(api_key=api_key)

def encode_image(image_path):
    """Encode the image to base64."""
    try:
        # 打开图片文件
        image = Image.open(image_path).convert("RGB")

        # 将图片转换为字节流
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

        return img_str
    except FileNotFoundError:
        print(f"Error: The file {image_path} was not found.")
        return None
    except Exception as e:  # 添加通用异常处理
        print(f"Error: {e}")
        return None

def feifeichat(image):
    model = "pixtral-large-2411"
    # Define the messages for the chat
    messages = [{
        "role":
        "user",
        "content": [
            {
                "type": "text",
                "text": "用英文详细描述下"
            },
            {
                "type": "image_url",
                "image_url": f"data:image/jpeg;base64,{base64_image}",
            },
        ],
    }]


    partial_message = Mistralclient.chat.stream(model=model, messages=messages):
    return partial_message


with gr.Blocks() as demo:
    gr.Markdown("Florence-2 Image To Flux Prompt")
    with gr.Tab(label="Image To Flux Prompt"):
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="Input Picture",height=480)
                submit_btn = gr.Button(value="Submit")
            with gr.Column():
                output_text = gr.Textbox(label="Output Text")


        submit_btn.click(feifeichat, [input_img], [output_text])

demo.launch()