app.py

import streamlit as st
import requests
import os
import base64
from PIL import Image
from io import BytesIO  # Correctly import BytesIO from the io module

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

# API key from environment variable
API_KEY = os.environ.get("NEBIUS_API_KEY")

if not API_KEY:
    st.error("API key not found. Please set the `NEBIUS_API_KEY` environment variable.")

# 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,
    }

    try:
        response = requests.post(api_url, json=payload, headers=headers)
        if response.status_code == 200:
            return response.json()
        else:
            st.error(f"API Error: {response.status_code}, {response.text}")
            return {"error": response.text}
    except Exception as e:
        st.error(f"An exception occurred: {e}")
        return {"error": str(e)}

# 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 = BytesIO()  # Use BytesIO from the io module
    image.save(buffered, format="PNG")
    image_base64 = base64.b64encode(buffered.getvalue()).decode()

    # Show the uploaded image
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Add a button to trigger caption generation
    if st.button("Generate Caption", use_container_width=True):
        st.write("Generating caption...")
        result = generate_caption(image_base64, API_KEY)

        # Extract and display the generated caption
        if "error" in result:
            st.error(f"Error: {result['error']}")
        else:
            try:
                # Extract the caption from the response
                caption = (
                    result.get("choices", [{}])[0]
                    .get("message", {})
                    .get("content", "No caption generated.")
                )
                st.subheader("Generated Caption")

                # Use a text area for the caption so it can be copied easily
                caption_area = st.text_area(
                    "", caption, height=100, key="caption_area", disabled=True,
                    max_chars=2000,  # Limit number of characters for better display
                    help="The generated caption will appear here.", label_visibility="collapsed"
                )

            except Exception as e:
                st.error(f"Error processing the response: {e}")
else:
    if not API_KEY:
        st.warning("Please set the `NEBIUS_API_KEY` environment variable.")
    else:
        st.info("Please upload an image.")

# Customizing CSS for the text area and button (with blue glowing effect for button)
st.markdown(
    """
    <style>
        /* Custom styles for the text area */
        .css-1w6p6e5 {
            background-color: #f0f0f5;
            color: #333333;
            font-size: 16px;
            padding: 10px;
        }
        
        /* Glowing blue button styling */
        .stButton>button {
            background-color: #007bff;  /* Blue color */
            color: white;
            font-size: 16px;
            border: none;
            border-radius: 4px;
            padding: 12px 24px;
            cursor: pointer;
            box-shadow: 0 4px 6px rgba(0, 123, 255, 0.6), 0 0 10px rgba(0, 123, 255, 0.8);
            transition: all 0.3s ease-in-out;
        }

        .stButton>button:hover {
            background-color: #0056b3;  /* Darker blue on hover */
            box-shadow: 0 4px 6px rgba(0, 123, 255, 0.8), 0 0 20px rgba(0, 123, 255, 1);
        }

        /* Text area styling */
        .stTextArea>div>textarea {
            background-color: #f5f5f5;
            color: #333;
            font-size: 14px;
            padding: 10px;
            border-radius: 5px;
        }

        .stTextArea>div>textarea:disabled {
            background-color: #e0e0e0;
            color: #666;
        }

        /* Centering elements */
        .css-1v3fvcr {
            justify-content: center;
            display: flex;
        }
        .css-ffhzg6 {
            text-align: center;
        }
    </style>
    """, unsafe_allow_html=True
)