feifeilib/feifeiflorence.py

from PIL import Image
from io import BytesIO
import base64
import requests
import os
from mistralai import Mistral
import gradio as gr

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


def encode_image(image_path):
    """Encode the image to base64."""
    try:
        # Open the image file
        image = Image.open(image_path).convert("RGB")

        # Resize the image to a height of 512 while maintaining the aspect ratio
        base_height = 512
        h_percent = base_height / float(image.size[1])
        w_size = int((float(image.size[0]) * float(h_percent)))
        image = image.resize((w_size, base_height), Image.LANCZOS)

        # Convert the image to a byte stream
        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:  # Add generic exception handling
        print(f"Error: {e}")
        return None


def feifeiflorence(image):
    try:
        model = "pixtral-12b-2409"
        # Define the messages for the chat
        base64_image = encode_image(image)
        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}",
                    },
                ],
                "stream": False,
            }
        ]

        partial_message = ""
        for chunk in Mistralclient.chat.stream(model=model, messages=messages):
            if chunk.data.choices[0].delta.content is not None:
                partial_message = partial_message + chunk.data.choices[0].delta.content
                yield partial_message
    except Exception as e:  # 添加通用异常处理
        print(f"Error: {e}")
        return "Please upload a photo"