Update app.py

app.py

@@ -3,35 +3,36 @@ import requests
 import os
 import base64
 from PIL import Image
-from io import BytesIO  # Correctly import BytesIO from the io module
+from io import BytesIO
 
 # Set page title and layout
-st.set_page_config(page_title="Image Caption Generator", layout="wide") 
+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
+# 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": "Qwen/Qwen2-VL-7B-Instruct",
         "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}"}},
-                ],
-            },
+            {"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."},
+                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
+            ]},
         ],
         "temperature": 0.6,
     }
@@ -49,58 +50,36 @@ def generate_caption(image_base64, api_key):
 
 # 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:
-    # Convert the image to base64
     image = Image.open(uploaded_image)
-    buffered = BytesIO()  # Use BytesIO from the io module
+    buffered = BytesIO()
     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 Prompt", use_container_width=True):
+    
+    if st.button("Generate Image", use_container_width=True):
         st.write("Generating Prompt...")
         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.")
-                )
+                caption = result.get("choices", [{}])[0].get("message", {}).get("content", "No caption generated.")
                 
-                st.text_area("Prompt Generated", caption, height=200, key="caption_output")
-
+                # 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_column_width=True)
+                
+                # Download button for the generated image
+                st.markdown(f"[Download Image]({image_url})", unsafe_allow_html=True)
+            
             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.")
-
-# Remove the sidebar and center the content
-st.markdown(
-    """
-    <style>
-        .css-1d391kg {text-align: center;}
-        .css-1v3fvcr {justify-content: center; display: flex;}
-        .css-1y4ccs5 {margin: 0 auto;}
-        .css-1aumxhk {display: flex; justify-content: center;}
-        .css-1k6kn8p {justify-content: center; align-items: center; display: flex;}
-        .css-ffhzg6 {text-align: center;}
-        textarea {
-            color: white !important;
-            background-color: #262626 !important;
-        }
-    </style>
-    """, unsafe_allow_html=True
-)
+    st.info("Please upload an image and select an aspect ratio.")