|
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" |
|
|
|
base64_image = encode_image(image_path) |
|
if not base64_image: |
|
return "Failed to encode the 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 += chunk.data.choices[0].delta.content |
|
return partial_message |
|
except Exception as e: |
|
print(f"Error: {e}") |
|
return "Please upload a valid 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, 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]) |
|
|
|
|
|
app = FastAPI() |
|
|
|
class PredictRequest(BaseModel): |
|
image_base64: str |
|
|
|
@app.post("/api/predict") |
|
def predict(request: PredictRequest): |
|
try: |
|
|
|
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) |
|
|
|
|
|
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) |
|
|