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

app = FastAPI()



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 feifeichat(image):
    try:
        model = "pixtral-large-2411"
        # 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"


with gr.Blocks(theme=theme, elem_id="app-container") as app:
    gr.Markdown("Image To Flux Prompt")
    with gr.Tab(label="Image To Prompt"):
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="Input Picture",height=320,type="filepath")
                submit_btn = gr.Button(value="Submit", variant='primary')
            with gr.Column():
                output_text = gr.Textbox(label="Flux Prompt", show_copy_button = True)
                clr_button =gr.Button("Clear",variant="primary", elem_id="clear_button")
                clr_button.click(lambda: gr.Textbox(value=""), None, output_text)

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

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