ImageCaption.py

import os
import base64
import gradio as gr
import requests
from typing import Literal

def encode_image_to_base64(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def generate_caption(image, caption_length: Literal["Short", "Long"]) -> str:
    try:
        # Get API key from environment variable
        api_key = os.getenv("OPENROUTER_API_KEY")
        if not api_key:
            raise ValueError("OpenRouter API key not found in environment variables")

        # Convert image to base64
        if isinstance(image, str):
            image_b64 = encode_image_to_base64(image)
        else:
            # Handle gradio file upload
            image_b64 = base64.b64encode(image.read()).decode('utf-8')

        # Prepare the API request
        headers = {
            "Authorization": f"Bearer {api_key}",
            "HTTP-Referer": os.getenv("YOUR_SITE_URL", "http://localhost:7860"),  # Default for local testing
            "X-Title": os.getenv("YOUR_APP_NAME", "Llama Image Captioner")
        }

        # Adjust the prompt based on caption length
        system_prompt = "You are an expert at describing images in detail."
        user_prompt = "Short and concise caption." if caption_length == "Short" else "Provide a detailed description of this image."

        payload = {
            "model": "meta-llama/llama-3.2-90b-vision-instruct",
            "messages": [
                {"role": "system", "content": system_prompt},
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": user_prompt},
                        {"type": "image_url", "image_url": f"data:image/jpeg;base64,{image_b64}"}
                    ]
                }
            ]
        }

        # Make the API request
        response = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=headers,
            json=payload
        )
        response.raise_for_status()

        # Extract and return the caption
        return response.json()["choices"][0]["message"]["content"]

    except Exception as e:
        return f"Error generating caption: {str(e)}"

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_caption,
    inputs=[
        gr.Image(type="filepath", label="Upload an image"),
        gr.Radio(["Short", "Long"], label="Caption Length", value="Short")
    ],
    outputs=gr.Textbox(label="Generated Caption"),
    title="Llama Image Captioner",
    description="Upload an image to generate a caption using Meta-Llama 3.2 90B Vision model.",
    examples=[],  # You can add example images here
    theme=gr.themes.Soft()
)

if __name__ == "__main__":
    iface.launch()