MegaTronX's picture
Rename app.py to app.bak
e3c61ec verified
import base64
import requests
import os
from mistralai import Mistral
import gradio as gr
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
api_key = os.getenv("MISTRAL_API_KEY")
Mistralclient = Mistral(api_key=api_key)
def encode_image(image_path):
"""Encode the image to base64."""
try:
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode('utf-8')
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_path):
try:
model = "pixtral-large-2411"
# Encode the input image to base64
base64_image = encode_image(image_path)
if not base64_image:
return "Failed to encode the 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}"
},
],
"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 += chunk.data.choices[0].delta.content
return partial_message
except Exception as e:
print(f"Error: {e}")
return "Please upload a valid photo."
# Gradio UI
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, type="filepath")
submit_btn = gr.Button(value="Submit")
with gr.Column():
output_text = gr.Textbox(label="Flux Prompt")
submit_btn.click(feifeichat, [input_img], [output_text])
# Add FastAPI for Remote Access
app = FastAPI()
class PredictRequest(BaseModel):
image_base64: str # Expecting a base64-encoded string
@app.post("/api/predict")
def predict(request: PredictRequest):
try:
# Decode the base64 image and save it temporarily
image_data = base64.b64decode(request.image_base64)
temp_image_path = "temp_image.jpg"
with open(temp_image_path, "wb") as temp_image:
temp_image.write(image_data)
# Run the prediction function
response = feifeichat(temp_image_path)
return {"prediction": response}
except Exception as e:
raise HTTPException(status_code=500, detail=f"Error during prediction: {e}")
demo.launch(app, share=True)