app.py

import gradio as gr
from PIL import Image
import requests
import os
from together import Together
import base64
import io

# Initialize Together client
client = None

def initialize_client(api_key=None):
    global client
    if api_key:
        os.environ["TOGETHER_API_KEY"] = api_key
    if "TOGETHER_API_KEY" in os.environ:
        client = Together()
    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,
        )

        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
    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)])
            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}
    #         ]
    #     })

    # Prepare the current message
    content = []
    user_text = ""

    try:
        content.append({
            "role": "user",
            "content": [
                {"type": "text", "text": message["text"]},
                {
                    "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

    # 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

    # Append the current user message to the API messages
    messages.append({"role": "user", "content": content})

    try:
        # Call the Together AI API with streaming
        stream = client.chat.completions.create(
            model="meta-llama/Llama-Vision-Free",
            messages=messages,
            max_tokens=max_new_tokens,
            temperature=temperature,
            stream=True,
        )

        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

    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])
        
        yield history

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!")
    
    with gr.Row():
        together_api_key = gr.Textbox(
            label="Together API Key",
            placeholder="Enter your TOGETHER_API_KEY here",
            type="password"
        )
    
    with gr.Row():
        max_new_tokens = gr.Slider(
            minimum=10,
            maximum=500,
            value=250,
            step=10,
            label="Maximum number of new tokens",
        )
        temperature = gr.Number(
            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(
        old_bot_streaming, 
        inputs=[msg, chatbot, together_api_key, max_new_tokens, temperature], 
        outputs=chatbot
    )
    clear.click(lambda: [], None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(debug=True)