Update app.py

app.py

@@ -12,6 +12,7 @@ import io
 from io import BytesIO
 import base64
 import re
+from mistralai import mistralai
 from gradio_client import Client
 from fake_useragent import UserAgent
 import random
@@ -20,6 +21,65 @@ 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"
+
 
 def query(prompt, negative_prompt, task, steps, sampler, cfg_scale, seed, width, height):
     result = {"prompt": prompt,"negative_prompt": negative_prompt,"task": task,"steps": steps,"sampler": sampler,"cfg_scale": cfg_scale,"seed": seed, "width": width, "height": height}
@@ -191,6 +251,19 @@ height: 100% !important;
 """
 
 with gr.Blocks(css=css, theme=theme) as app:
+    with gr.Tab(label="Image To Prompt"):
+        with gr.Row():
+            with gr.Column(scale=4, min_width=300):
+                input_img = gr.Image(label="Input Picture", type="filepath")
+                
+            with gr.Column(scale=3):
+                output_text = gr.Textbox(label="Flux Prompt", lines=2, scale=6, show_copy_button = True)
+                submit_btn = gr.Button(value="Generate Pompt", scale=4, variant='primary')
+                clear_prompt =gr.Button("Clear 🗑️",variant="primary", elem_id="clear_button")
+                clear_prompt.click(lambda: (None, None), None, [input_img, output_text], queue=False, show_api=False)
+    
+    submit_btn.click(feifeichat, [input_img], [output_text])
+
 
     with gr.Tab("Basic Settings"):
         with gr.Row():