app.py

import gradio as gr
from groq import Groq
import base64
import os
import io
import json
from PIL import Image
import traceback

# Custom CSS for styling
custom_css = """
.center-aligned {
    text-align: center !important;
    color: #ff4081;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
}
.input-background {
    background-color: #B7E0FF !important;
    padding: 15px !important;
    border-radius: 10px !important;
    margin: 0 !important;
}
.input-background textarea {
    font-size: 18px !important;
    background-color: #ffffff;
    border: 1px solid #f0f8ff;
    border-radius: 8px;
}
.image-background {
    border-radius: 10px !important;
    border: 2px solid #B7E0FF !important;
}
.lng-background {
    background-color: #FFF5CD !important;
    padding: 5px !important;
    border-radius: 10px !important;
    margin: 0 !important;
}
.api-background {
    background-color: #FFCFB3 !important;
    padding: 5px !important;
    border-radius: 10px !important;
    margin: 0 !important;
}
.script-background {
    background-color: #FEF9D9 !important;
    padding: 15px !important;
    border-radius: 10px !important;
    margin: 0 !important;
}
.script-background textarea {
    font-size: 18px !important;
    background-color: #ffffff;
    border: 1px solid #f0f8ff;
    border-radius: 8px;
}
.model-background {
    background-color: #FFF4B5 !important;
    padding: 5px !important;
    border-radius: 10px !important;
    margin: 0 !important;
}
.gen-button {
    border-radius: 10px !important;
    border: none !important;
    background-color: #ff4081 !important;
    color: white !important;
    font-weight: bold !important;
    transition: all 0.3s ease !important;
    margin: 0 !important;
}
.gen-button:hover {
    background-color: #f50057 !important;
    transform: scale(1.05);
}
.clear-button {
    color: white !important;
    background-color: #000000 !important;
    padding: 5px !important;
    border-radius: 10px !important;
    margin: 0 !important;
}
.clear-button:hover {
    background-color: #000000 !important;
    transform: scale(1.05);
}
"""

# List of available models
MODELS = [
    "llama-3.2-90b-vision-preview",
    "llama-3.2-11b-vision-preview",
    "llava-v1.5-7b-4096-preview"
]

def compress_image(image, max_size=(800, 800), quality=95):
    img = Image.open(image) if isinstance(image, str) else image
    img.thumbnail(max_size)
    buffered = io.BytesIO()
    img.save(buffered, format="JPEG", quality=quality)
    return buffered.getvalue()

def encode_image(image):
    if isinstance(image, Image.Image):
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG", quality=95)
        return base64.b64encode(buffered.getvalue()).decode('utf-8')
    else:
        compressed = compress_image(image)
        return base64.b64encode(compressed).decode('utf-8')

def create_client():
    api_key = "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"  # Directly using API key in the code - Hardcoding it for a purpose. I know I should keep it separate but I am not doing it on purpose<3.
    return Groq(api_key=api_key)

def analyze_input(text_input, Quick_Input):
    if Quick_Input == "Input Manually":
        return text_input.strip()
    elif Quick_Input == "Describe Image":
        return "Take a close look at the image and describe it in as much detail as possible. Be sure to mention the main subject, the background, the colors used, the mood or feeling it evokes, and any specific elements that stand out."
    elif Quick_Input == "Text in Image":
        return "What does the text in this photo say?"
    elif Quick_Input == "Image Reasoning":
        return "Let's work this out in a step by step way to be sure we have the right answer. Deduce from the image and provide a quick answer."
    else:
        return text_input.strip()

def process_image_and_text(image, text_input, Quick_Input, model):
    gr.Info("Image is being analyzed, please wait a moment...")
    if Quick_Input == "Input Manually" and not text_input.strip():
        return "Error: Please enter a question or select a quick input option!"
    
    text_input = text_input.strip()
    client = create_client()
    base64_image = encode_image(image)
    Input_Text = analyze_input(text_input, Quick_Input)
    
    try:
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": Input_Text},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}",
                            },
                        },
                    ],
                }
            ],
            model=model,
            temperature=1,
        )
        gr.Info("Vision model has completed the response.")
        response_content = chat_completion.choices[0].message.content.strip()
        return response_content  # No language conversion needed, only in English.
    except Exception as e:
        error_traceback = traceback.format_exc()
        print(f"Error occurred: {error_traceback}")
        return f"Error occurred: {str(e)}"

def update_textbox_based_on_quick_input(Quick_Input):
    return analyze_input("", Quick_Input)

with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css) as iface:
    gr.Markdown("""
    # Image Analysis and Reasoning - Large Vision Model
    > ### **Project by Muhammad John Abbas - Supervised by: Dr. Mudassar Raza**
    """, elem_classes="center-aligned")

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Upload Image", elem_classes="image-background")
            model_select = gr.Dropdown(choices=MODELS, label="Select Vision Model", value=MODELS[0], elem_classes="model-background")
            clear_button = gr.Button("Clear Answer and Image", variant="secondary", elem_classes="clear-button")
            gr.Markdown("""
            ### **※ Can analyze image and search for specific objects and events**
            """, elem_classes="center-aligned")

        with gr.Column(scale=1):
            text_input = gr.Textbox(label="Enter Question", placeholder="Please enter your question...", autofocus=True, elem_classes="input-background", max_lines=5)
            with gr.Row():
                Quick_Input = gr.Dropdown(
                    choices=["Input Manually", "Describe Image", "Image Reasoning", "Text in Image"],
                    value="Input Manually",
                    label="Quick Input",
                    interactive=True,
                    elem_classes="lng-background"
                )
            submit_button = gr.Button("Submit", variant="primary", elem_classes="gen-button")
            output = gr.Textbox(label="Vision Model Response", elem_classes="script-background", max_lines=40)

    Quick_Input.change(
        fn=update_textbox_based_on_quick_input,
        inputs=[Quick_Input],
        outputs=[text_input]
    )

    submit_button.click(
        fn=process_image_and_text,
        inputs=[image_input, text_input, Quick_Input, model_select],
        outputs=[output]
    )

    def clear_outputs():
        return None, None, ""

    clear_button.click(
        fn=clear_outputs,
        inputs=[],
        outputs=[image_input, text_input, output]
    )

if __name__ == "__main__":
    iface.launch(share=True, show_api=False)