import streamlit as st
import requests
import os
import base64
from PIL import Image

# Set page title
st.set_page_config(page_title="Image Caption Generator", layout="centered")

# UI for the app
st.title("Image Caption Generator")
st.write(
    "Upload an image, and this app will generate a detailed caption for it using the Nebius AI API."
)

# Sidebar for API key
api_key = st.sidebar.text_input(
    "Enter Nebius API Key", type="password", help="Add your Nebius API key here."
)

# Function to call Nebius API
def generate_caption(image_base64, api_key):
    api_url = "https://api.studio.nebius.ai/v1/chat/completions"
    headers = {"Authorization": f"Bearer {api_key}"}
    payload = {
        "model": "Qwen/Qwen2-VL-72B-Instruct",
        "messages": [
            {
                "role": "system",
                "content": """You are an image to prompt converter. Your work is to observe each and every detail of the image and craft a detailed prompt under 75 words in this format: [image content/subject, description of action, state, and mood], [art form, style], [artist/photographer reference if needed], [additional settings such as camera and lens settings, lighting, colors, effects, texture, background, rendering].""",
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Write a caption for this image"},
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
                ],
            },
        ],
        "temperature": 0,
    }

    response = requests.post(api_url, json=payload, headers=headers)
    if response.status_code == 200:
        return response.json()
    else:
        return {"error": response.text}

# File uploader for image
uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

if uploaded_image and api_key:
    # Convert the image to base64
    image = Image.open(uploaded_image)
    buffered = st.BytesIO()
    image.save(buffered, format="PNG")
    image_base64 = base64.b64encode(buffered.getvalue()).decode()

    # Generate caption
    st.image(image, caption="Uploaded Image", use_column_width=True)
    st.write("Generating caption...")
    result = generate_caption(image_base64, api_key)

    # Display the result
    if "error" in result:
        st.error(f"Error: {result['error']}")
    else:
        caption = result.get("messages", [{}])[-1].get("content", [{}])[0].get("text", "No caption generated.")
        st.subheader("Generated Caption")
        st.write(caption)
else:
    st.info("Please upload an image and provide your API key.")

st.sidebar.write("Built with ❤️ by OpenAI GPT-4")