Update app.py

app.py

@@ -1,83 +1,107 @@
 import streamlit as st
 import requests
+import base64
 import os
-import json
 
-# Set the Nebius API key
-API_KEY = os.environ.get("NEBIUS_API_KEY")
-API_URL = "https://api.studio.nebius.ai/v1/chat/completions"
+# Function to convert image to base64
+def convert_image_to_base64(image):
+    image_bytes = image.read()
+    encoded_image = base64.b64encode(image_bytes).decode("utf-8")
+    return encoded_image
 
-# Streamlit app configuration
-st.set_page_config(page_title="Image to Prompt Converter", layout="centered", page_icon="🖼️")
+# Function to generate a caption using Nebius API
+def generate_caption(encoded_image):
+    API_URL = "https://api.studio.nebius.ai/v1/chat/completions"
+    API_KEY = os.environ.get("NEBIUS_API_KEY")
 
-# Apply custom styles
-st.markdown(
-    """
-    <style>
-    body {
-        background: linear-gradient(135deg, #1e3c72, #2a5298);
-        color: #FFFFFF;
-        font-family: 'Arial', sans-serif;
+    headers = {
+        "Authorization": f"Bearer {API_KEY}",
+        "Content-Type": "application/json"
     }
-    .stApp {
-        align-items: center;
-        justify-content: center;
-    }
-    img {
-        border-radius: 10px;
-        margin-bottom: 20px;
-        max-width: 100%;
-        height: auto;
-    }
-    h1, h2, h3 {
-        text-align: center;
-    }
-    </style>
-    """,
-    unsafe_allow_html=True
-)
 
-# App title
-st.title("Image to Prompt Converter")
-st.markdown("**Upload an image and generate a detailed prompt.**")
+    payload = {
+        "model": "llava-hf/llava-1.5-7b-hf",
+        "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": "write a detailed caption for this image"
+            }
+        ],
+        "image": {
+            "type": "image_url",
+            "image_url": {
+                "url": f"data:image/png;base64,{encoded_image}"
+            }
+        },
+        "temperature": 0.7
+    }
 
-# Image upload
-uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+    response = requests.post(API_URL, headers=headers, json=payload)
+    
+    if response.status_code == 200:
+        result = response.json()
+        caption = result.get("choices", [{}])[0].get("message", {}).get("content", "No caption generated.")
+        return caption
+    else:
+        st.error(f"API Error {response.status_code}: {response.text}")
+        return None
 
-if uploaded_file is not None:
-    # Display the uploaded image
-    st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
+# Streamlit app
+def main():
+    st.set_page_config(page_title="Image to Caption Converter", layout="centered", initial_sidebar_state="collapsed")
     
-    # Generate button
-    if st.button("Generate Prompt"):
-        # Prepare the API payload
-        files = {"file": uploaded_file.getvalue()}
-        headers = {"Authorization": f"Bearer {API_KEY}"}
-        data = {
-            "model": "Qwen/Qwen2-VL-7B-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]."""
-                }
-            ],
-            "temperature": 1
-        }
+    # Gradient background style
+    st.markdown("""
+        <style>
+            body {
+                background: linear-gradient(135deg, #1e3c72, #2a5298);
+                color: white;
+                font-family: 'Arial', sans-serif;
+            }
+            .uploaded-image {
+                max-width: 100%;
+                border: 2px solid #ffffff;
+                border-radius: 10px;
+            }
+            .copy-button {
+                background-color: #ff8800;
+                color: white;
+                border: none;
+                border-radius: 5px;
+                padding: 10px 15px;
+                cursor: pointer;
+            }
+            .copy-button:hover {
+                background-color: #cc6b00;
+            }
+        </style>
+    """, unsafe_allow_html=True)
+
+    st.title("🖼️ Image to Caption Converter")
+
+    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+    if uploaded_file:
+        # Display the uploaded image
+        st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
+
+        # Convert image to base64 and get caption
+        if st.button("Generate Caption"):
+            with st.spinner("Generating caption..."):
+                encoded_image = convert_image_to_base64(uploaded_file)
+                caption = generate_caption(encoded_image)
+
+                if caption:
+                    st.subheader("Generated Caption:")
+                    st.text_area("", caption, height=100, key="caption_area")
+
+                    # Copy button
+                    if st.button("Copy to Clipboard"):
+                        st.code(caption, language="text")
+                        st.success("Caption copied to clipboard!")
 
-        # Call the Nebius API
-        response = requests.post(API_URL, headers=headers, data=json.dumps(data))
-        
-        if response.status_code == 200:
-            # Extract the generated prompt
-            result = response.json()
-            generated_prompt = result.get("choices", [{}])[0].get("message", {}).get("content", "No prompt generated.")
-            
-            # Display the generated prompt
-            st.subheader("Generated Prompt")
-            st.text_area("", generated_prompt, height=200)
-            
-            # Copy button
-            if st.button("Copy Prompt"):
-                st.write("Copy functionality is not supported in this environment. Please manually copy the text.")
-        else:
-            st.error(f"Failed to generate prompt: {response.status_code} - {response.text}")
+if __name__ == "__main__":
+    main()