Spaces:
Running
Running
| 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 | |
| st.set_page_config(page_title="Image Caption Generator", layout="centered") | |
| # API key from environment variable | |
| API_KEY = os.environ.get("NEBIUS_API_KEY") | |
| # 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." | |
| ) | |
| 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, | |
| } | |
| 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 = 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_column_width=True) | |
| # Add a button to trigger caption generation | |
| if st.button("Generate Caption"): | |
| 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: | |
| if not API_KEY: | |
| st.warning("Please set the `NEBIUS_API_KEY` environment variable.") | |
| else: | |
| st.info("Please upload an image.") | |
| st.sidebar.write("Built with ❤️ by OpenAI GPT-4") |