Update app.py

app.py

@@ -1,62 +1,54 @@
+import base64
+from io import BytesIO
+import os
+from mistralai import Mistral
+import re
 from PIL import Image
-import gradio as gr
-from transformers import (
-    AutoProcessor,
-    AutoModelForCausalLM,
-)
-import torch
-import subprocess
 
-subprocess.run(
-    "pip install flash-attn --no-build-isolation",
-    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-    shell=True,
-)
-
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-Florence_models = AutoModelForCausalLM.from_pretrained(
-    "microsoft/Florence-2-large",
-    torch_dtype=torch_dtype,
-    trust_remote_code=True).to(device)
-
-Florence_processors = AutoProcessor.from_pretrained(
-    "microsoft/Florence-2-large", trust_remote_code=True)
-
-
-
-def feifeiflorence(
-        image,
-        progress=gr.Progress(track_tqdm=True),
-):
-    image = Image.fromarray(image)
-    task_prompt = "<MORE_DETAILED_CAPTION>"
-
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-
-    inputs = Florence_processors(text=task_prompt,
-                                 images=image,
-                                 return_tensors="pt").to(device, torch_dtype)
-
-    generated_ids = Florence_models.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
-        max_new_tokens=1024,
-        num_beams=3,
-        do_sample=False,
-    )
-    generated_text = Florence_processors.batch_decode(
-        generated_ids, skip_special_tokens=False)[0]
-    parsed_answer = Florence_processors.post_process_generation(
-        generated_text,
-        task=task_prompt,
-        image_size=(image.width, image.height))
-    out_text=parsed_answer["<MORE_DETAILED_CAPTION>"]
-    width, height = image.size
-    return out_text,f"width={width} height={height}"
 
+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):
+    model = "pixtral-large-2411"
+    # Define the messages for the chat
+    messages = [{
+        "role":
+        "user",
+        "content": [
+            {
+                "type": "text",
+                "text": "用英文详细描述下"
+            },
+            {
+                "type": "image_url",
+                "image_url": f"data:image/jpeg;base64,{base64_image}",
+            },
+        ],
+    }]
+
+
+    partial_message = Mistralclient.chat.stream(model=model, messages=messages):
+    return partial_message
 
 
 with gr.Blocks() as demo:
@@ -68,9 +60,8 @@ with gr.Blocks() as demo:
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():
                 output_text = gr.Textbox(label="Output Text")
-                output_img_text = gr.Textbox(label="Output width and height")
 
 
-        submit_btn.click(feifeiflorence, [input_img], [output_text, output_img_text])
+        submit_btn.click(feifeichat, [input_img], [output_text])
 
 demo.launch()