app.py

import os
import gradio as gr
from argparse import ArgumentParser
from groq import Groq
from PIL import Image
import base64
import io

# Initialize Groq client
API_KEY = os.environ['GROQ_API_KEY']
client = Groq(api_key=API_KEY)

REVISION = 'v1.0.4'

def _get_args():
    parser = ArgumentParser()
    parser.add_argument("--revision", type=str, default=REVISION)
    parser.add_argument("--share", action="store_true", default=False, help="Create a publicly shareable link for the interface.")
    return parser.parse_args()

def process_image(image):
    # Convert image to bytes for Groq API
    buffered = io.BytesIO()
    image.save(buffered, format="JPEG")
    return buffered.getvalue()

def create_messages(query, image_data):
    messages = []
    
    # User query as text
    if query:
        messages.append({'role': 'user', 'content': query})

    # Include image if provided
    if image_data:
        image_base64 = f"data:image/jpeg;base64,{base64.b64encode(image_data).decode()}"
        messages.append({
            'role': 'user',
            'content': [
                {"type": "text", "text": "Please analyze this image."},
                {"type": "image_url", "image_url": {"url": image_base64}}
            ]
        })
    
    return messages

def predict(chat_history, query, image):
    # Process the image if provided
    image_data = process_image(image) if image else None
    messages = create_messages(query, image_data)

    # Call the Groq API with the messages
    try:
        completion = client.chat.completions.create(
            model="llama-3.2-11b-vision-preview",
            messages=messages,
            temperature=1,
            max_tokens=1500,
            top_p=1,
            stream=False,
        )

        response_text = completion.choices[0].message.content.strip()
    except Exception as e:
        response_text = f"Error: {str(e)}"
    
    chat_history.append((query, response_text))
    return chat_history

def clear_history():
    return []

def main():
    args = _get_args()
    
    with gr.Blocks() as demo:
        gr.Markdown("<h1 style='text-align: center;'>Llama-3.2-11b-vision-preview</h1>")
        
        chatbox = gr.Chatbot()
        query = gr.Textbox(label="Input", placeholder="Type your query here...")
        image_input = gr.Image(type="pil", label="Upload Image")
        
        submit_btn = gr.Button("Submit")
        clear_btn = gr.Button("Clear History")
        
        submit_btn.click(predict, inputs=[chatbox, query, image_input], outputs=chatbox)
        clear_btn.click(clear_history, outputs=chatbox)
        
    demo.launch(share=args.share)

if __name__ == '__main__':
    main()