Spaces:
Running
Running
| import streamlit as st | |
| import requests | |
| import os | |
| import base64 | |
| from PIL import Image | |
| from io import BytesIO | |
| # Set page title and layout | |
| st.set_page_config(page_title="AI Image Generator", layout="wide") | |
| # Aspect Ratio Selection | |
| aspect_ratios = { | |
| "16:9 (1280x720)": (1280, 720), | |
| "1:1 (1080x1080)": (1080, 1080), | |
| "9:16 (720x1280)": (720, 1280), | |
| "2:3 (800x1200)": (800, 1200), | |
| } | |
| # 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 for prompt generation | |
| 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": "mistralai/Mistral-Small-3.1-24B-Instruct-2503", | |
| "messages": [ | |
| {"role": "system", "content": "You are an image-to-text prompt converter for AI image generation."}, | |
| {"role": "user", "content": [ | |
| {"type": "text", "text": "Describe This Image In Detail under 100 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]."}, | |
| {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}}, | |
| ]}, | |
| ], | |
| "temperature": 0.6, | |
| } | |
| 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"]) | |
| selected_aspect = st.selectbox("Select Aspect Ratio", list(aspect_ratios.keys())) | |
| width, height = aspect_ratios[selected_aspect] | |
| if uploaded_image and API_KEY: | |
| image = Image.open(uploaded_image) | |
| buffered = BytesIO() | |
| image.save(buffered, format="PNG") | |
| image_base64 = base64.b64encode(buffered.getvalue()).decode() | |
| st.image(image, caption="Uploaded Image", use_container_width=True) | |
| if st.button("Generate Image", use_container_width=True): | |
| st.write("Wait AI Is Generating Image...") | |
| result = generate_caption(image_base64, API_KEY) | |
| if "error" in result: | |
| st.error(f"Error: {result['error']}") | |
| else: | |
| try: | |
| caption = result.get("choices", [{}])[0].get("message", {}).get("content", "No caption generated.") | |
| # Construct the image generation URL | |
| image_url = f"https://image.pollinations.ai/prompt/{requests.utils.quote(caption)}?width={width}&height={height}&nologo=true&model=flux&enhance=true&seed=119" | |
| st.image(image_url, caption="Generated Image", use_container_width=True) | |
| # Download button for the generated image | |
| st.download_button( | |
| label="Download Image", | |
| data=requests.get(image_url).content, # Fetch the image data | |
| file_name="generated_image.png", | |
| mime="image/png", | |
| use_container_width=True | |
| ) | |
| except Exception as e: | |
| st.error(f"Error processing the response: {e}") | |
| else: | |
| st.info("Please upload an image and select an aspect ratio.") | |