Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from PIL import Image
-import requests
 import os
 from together import Together
 import base64
@@ -9,6 +8,7 @@ import io
 # Initialize Together client
 client = None
 
+
 def initialize_client(api_key=None):
     global client
     if api_key:
@@ -18,174 +18,81 @@ def initialize_client(api_key=None):
     else:
         raise ValueError("Please provide a Together API Key")
 
-def encode_image(image_path):
-    try:
-        with Image.open(image_path) as img:
-            buffered = io.BytesIO()
-            img.save(buffered, format="PNG")
-            return base64.b64encode(buffered.getvalue()).decode('utf-8')
-    except Exception as e:
-        print(f"Error encoding image: {e}")
-        raise e
-
-def old_bot_streaming(message, history, max_new_tokens=250, api_key=None, max_history=5):
-    if client is None:
-        initialize_client(api_key)
-
-    txt = message["text"]
-    messages = []
-    images = []
-
-    for i, msg in enumerate(history[-max_history:]):
-        if isinstance(msg[0], tuple):
-            messages.append({"role": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(msg[0][0])}"}}]})
-            messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
-        elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
-            pass
-        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str):
-            messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
-            messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})
 
-    if len(message["files"]) == 1:
-        if isinstance(message["files"][0], str):  # examples
-            image_path = message["files"][0]
-        else:  # regular input
-            image_path = message["files"][0]["path"]
-        messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path)}"}}]})
-    else:
-        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
-
-    try:
-        stream = client.chat.completions.create(
-            model="meta-llama/Llama-Vision-Free",
-            messages=messages,
-            max_tokens=max_new_tokens,
-            stream=True,
-        )
+def encode_image(image_path):
+    with Image.open(image_path) as img:
+        buffered = io.BytesIO()
+        img.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-        buffer = ""
-        for chunk in stream:
-            if chunk.choices[0].delta.content is not None:
-                buffer += chunk.choices[0].delta.content
-                time.sleep(0.01)
-                yield buffer
 
-    except together.error.InvalidRequestError as e:
-        if "Request Entity Too Large" in str(e):
-            yield "The image is too large. Please try with a smaller image or compress the existing one."
-        else:
-            yield f"An error occurred: {str(e)}"
-            
-def bot_streaming(message, history, together_api_key, max_new_tokens=250, temperature=0.7):
-    # Initialize history if it's None
-    if history is None:
-        history = []
-    
-    # Initialize the Together client if not already done
+def bot_streaming(
+    message, history, together_api_key, max_new_tokens=250, temperature=0.7
+):
     if client is None:
         try:
             initialize_client(together_api_key)
         except Exception as e:
-            # Append error to history and yield
-            history.append(["Error initializing client", str(e)])
+            history.append((message, f"Error initializing client: {str(e)}"))
             yield history
             return
 
     prompt = "You are a helpful AI assistant. Analyze the image provided (if any) and respond to the user's query or comment."
-
     messages = [{"role": "system", "content": prompt}]
 
-    # # Build the conversation history for the API
-    # for idx, (user_msg, assistant_msg) in enumerate(history):
-    #     # Append user messages
-    #     messages.append({
-    #         "role": "user",
-    #         "content": [
-    #             {"type": "text", "text": user_msg}
-    #         ]
-    #     })
-    #     # Append assistant messages
-    #     messages.append({
-    #         "role": "assistant",
-    #         "content": [
-    #             {"type": "text", "text": assistant_msg}
-    #         ]
-    #     })
+    # Add history to messages
+    for user_msg, assistant_msg in history:
+        if isinstance(user_msg, str):  # Text message
+            messages.append(
+                {"role": "user", "content": [{"type": "text", "text": user_msg}]}
+            )
+        elif isinstance(user_msg, dict):  # Image message
+            image_base64 = encode_image(user_msg["image_path"])
+            messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": user_msg.get("text", "")},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/png;base64,{image_base64}"
+                            },
+                        },
+                    ],
+                }
+            )
+        messages.append(
+            {"role": "assistant", "content": [{"type": "text", "text": assistant_msg}]}
+        )
 
     # Prepare the current message
-    content = []
-    user_text = ""
-
-    try:
-        image_path = ''
-        if 'files' in message and len(message['files']) > 0:
-            file_info = message['files'][0]
-            if isinstance(file_info, dict) and 'path' in file_info:
-                image_path = file_info['path']
-            elif isinstance(file_info, str):
-                image_path = file_info
-            else:
-                raise ValueError("Invalid file information provided.")
-        
-        content.append({
-            "role": "user",
-            "content": [
-                {"type": "text", "text": message["text"]},
+    user_message_content = []
+    if isinstance(message, dict):
+        if message.get("text"):
+            user_message_content.append({"type": "text", "text": message["text"]})
+        if message.get("files") and len(message["files"]) > 0:
+            image_path = message["files"][0]
+            image_base64 = encode_image(image_path)
+            user_message_content.append(
                 {
                     "type": "image_url",
-                    "image_url": {
-                        "url": f"data:image/png;base64,{encode_image(image_path)}",
-                    },
-                },
-            ],
-        })
-
-        # if isinstance(message, dict):
-        #     # Handle text input
-        #     if 'text' in message and message['text']:
-        #         user_text = message['text']
-        #         content.append({"type": "text", "text": user_text})
-            
-        #     # Handle image input
-        #     if 'files' in message and len(message['files']) > 0:
-        #         file_info = message['files'][0]
-        #         if isinstance(file_info, dict) and 'path' in file_info:
-        #             image_path = file_info['path']
-        #         elif isinstance(file_info, str):
-        #             image_path = file_info
-        #         else:
-        #             raise ValueError("Invalid file information provided.")
-
-        #         # Encode the image to base64
-        #         image_base64 = encode_image(image_path)
-        #         content.append({
-        #             "type": "image_url",
-        #             "image_url": {"url": f"data:image/png;base64,{image_base64}"}
-        #         })
-        #         user_text += "\n[User uploaded an image]"
-        # else:
-        #     # If message is a string
-        #     user_text = message
-        #     content.append({"type": "text", "text": user_text})
-    except Exception as e:
-        # If there's an error processing the input, append it to history and yield
-        error_message = f"An error occurred while processing your input: {str(e)}"
-        print(error_message)  # Debug statement
-        history.append([user_text or "[Invalid input]", error_message])
-        yield history
-        return
+                    "image_url": {"url": f"data:image/png;base64,{image_base64}"},
+                }
+            )
+    elif isinstance(message, str):
+        user_message_content.append({"type": "text", "text": message})
 
-    # Append the new user message with an empty assistant response
-    history.append([user_text, ""])
-    yield history  # Yield the updated history to show the user's message immediately
+    current_message = {"role": "user", "content": user_message_content}
+    messages.append(current_message)
 
-    # Append the current user message to the API messages
-    messages.append({"role": "user", "content": content})
+    # Add the user's message to the history
+    user_display_message = message["text"] if isinstance(message, dict) else message
+    history = history + [(user_display_message, "")]
 
     try:
-        # Call the Together AI API with streaming
         stream = client.chat.completions.create(
-            model="meta-llama/Llama-Vision-Free",
+            model="meta-llama/Llama-3.2-11B-Vision-Instruct-Turbo",
             messages=messages,
             max_tokens=max_new_tokens,
             temperature=temperature,
@@ -194,52 +101,47 @@ def bot_streaming(message, history, together_api_key, max_new_tokens=250, temper
 
         response = ""
         for chunk in stream:
-            # Extract the content from the API response
-            chunk_content = chunk.choices[0].delta.content or ""
-            print(chunk.choices[0].delta.content or "", end="", flush=True)
-            response += chunk_content
-            # # Update the last assistant message in history
-            # if history:
-            #     history[-1][1] = response
-            #     yield history
-            # else:
-            #     # If history is somehow empty, append the response
-            #     history.append(["", response])
-            #     yield history
-
-        # if not response:
-        #     # If no response was generated, notify the user
-        #     history[-1][1] = "No response generated. Please try again."
-        #     yield history
+            if (
+                chunk.choices
+                and chunk.choices[0].delta
+                and chunk.choices[0].delta.content is not None
+            ):
+                response += chunk.choices[0].delta.content
+                # Update the assistant's response in the history
+                history[-1] = (user_display_message, response)
+                yield history
+
+        if not response:
+            history[-1] = (
+                user_display_message,
+                "No response generated. Please try again.",
+            )
+            yield history
 
     except Exception as e:
-        # Handle exceptions from the API call
-        error_message = ""
-        if "Request Entity Too Large" in str(e):
-            error_message = "The image is too large. Please try with a smaller image or compress the existing one."
-        else:
-            error_message = f"An error occurred: {str(e)}"
-        
-        print(error_message)  # Debug statement
-
-        if history:
-            history[-1][1] = error_message
-        else:
-            history.append(["", error_message])
-        
+        error_message = (
+            "The image is too large. Please try with a smaller image or compress the existing one."
+            if "Request Entity Too Large" in str(e)
+            else f"An error occurred: {str(e)}"
+        )
+        history[-1] = (user_display_message, error_message)
         yield history
 
+
+# The rest of your Gradio interface code remains the same
 with gr.Blocks() as demo:
     gr.Markdown("# Meta Llama-3.2-11B-Vision-Instruct (FREE)")
-    gr.Markdown("Try the new Llama 3.2 11B Vision API by Meta for free through Together AI. Upload an image, and start chatting about it. Just paste in your Together AI API key and get started!")
-    
+    gr.Markdown(
+        "Try the new Llama 3.2 11B Vision API by Meta for free through Together AI. Upload an image, and start chatting about it. Just paste in your Together AI API key and get started!"
+    )
+
     with gr.Row():
         together_api_key = gr.Textbox(
             label="Together API Key",
             placeholder="Enter your TOGETHER_API_KEY here",
-            type="password"
+            type="password",
         )
-    
+
     with gr.Row():
         max_new_tokens = gr.Slider(
             minimum=10,
@@ -249,23 +151,19 @@ with gr.Blocks() as demo:
             label="Maximum number of new tokens",
         )
         temperature = gr.Number(
-            value=0.7,
-            minimum=0,
-            maximum=1,
-            step=0.1,
-            label="Temperature"
+            value=0.7, minimum=0, maximum=1, step=0.1, label="Temperature"
         )
-    
+
     chatbot = gr.Chatbot()
     msg = gr.MultimodalTextbox(label="Enter text or upload an image")
     clear = gr.Button("Clear")
 
     msg.submit(
-        bot_streaming, 
-        inputs=[msg, chatbot, together_api_key, max_new_tokens, temperature], 
-        outputs=chatbot
+        bot_streaming,
+        [msg, chatbot, together_api_key, max_new_tokens, temperature],
+        chatbot,
     )
-    clear.click(lambda: [], None, chatbot, queue=False)
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
     demo.launch(debug=True)