app.py

import gradio as gr
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):
    try:
        model = "pixtral-large-2411"
        base64_image = encode_image(image)
        # Define the messages for the chat
        messages = [{
            "role":
            "user",
            "content": [
                {
                    "type": "text",
                    "text": "Please provide a detailed description of this photo"
                },
                {
                    "type": "image_url",
                    "image_url": f"data:image/jpeg;base64,{base64_image}",
                },
            ],
        }]
    
    
        partial_message = Mistralclient.chat.stream(model=model, messages=messages)
        return partial_message
    except Exception as e:  # 添加通用异常处理
        print(f"Error: {e}")
        return "Please upload a photo"


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()