Update app.py

app.py

@@ -1,69 +1,74 @@
-import os
 import streamlit as st
-from openai import OpenAI
+import requests
+import os
+import base64
 from PIL import Image
-import io
 
-# Set up the OpenAI client
-client = OpenAI(
-    base_url="https://api.studio.nebius.ai/v1/",
-    api_key=os.environ.get("NEBIUS_API_KEY")
+# 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."
 )
 
-# Function to generate caption from image URL
-def generate_caption(image_data):
-    completion = client.chat.completions.create(
-        model="Qwen/Qwen2-VL-72B-Instruct",
-        messages=[
+# 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]."""
+                "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": image_data
-                        }
-                    }
-                ]
-            }
+                    {"type": "text", "text": "Write a caption for this image"},
+                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
+                ],
+            },
         ],
-        temperature=0
-    )
-    
-    caption = completion.to_json().get("choices", [{}])[0].get("message", {}).get("content", "")
-    return caption
+        "temperature": 0,
+    }
 
-# Streamlit UI
-st.title("Image to Caption Generator")
-st.write("Upload an image, and the app will generate a detailed caption for it.")
+    response = requests.post(api_url, json=payload, headers=headers)
+    if response.status_code == 200:
+        return response.json()
+    else:
+        return {"error": response.text}
 
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+# File uploader for image
+uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
 
-if uploaded_file is not None:
-    # Display the uploaded image
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    
-    # Convert image to a base64 string
-    buffered = io.BytesIO()
+if uploaded_image and api_key:
+    # Convert the image to base64
+    image = Image.open(uploaded_image)
+    buffered = st.BytesIO()
     image.save(buffered, format="PNG")
-    img_base64 = buffered.getvalue().decode("utf-8")
+    image_base64 = base64.b64encode(buffered.getvalue()).decode()
 
-    # Generate caption using the OpenAI API
+    # Generate caption
+    st.image(image, caption="Uploaded Image", use_column_width=True)
     st.write("Generating caption...")
-    caption = generate_caption(img_base64)
+    result = generate_caption(image_base64, api_key)
 
-    # Display the generated caption
-    if caption:
-        st.subheader("Generated Caption:")
-        st.write(caption)
+    # Display the result
+    if "error" in result:
+        st.error(f"Error: {result['error']}")
     else:
-        st.write("No caption could be generated.")
+        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")