app.py

import streamlit as st
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import base64

# Initialize Gemini API client
client = genai.Client(api_key=st.secrets["GEMINI_API_KEY"])

st.title("RINGGEN: AI-Powered Engagement Ring Designer")

# Dropdown and input fields for customization
ring_style = st.selectbox("Ring Style", ["Solitaire", "Halo", "Hidden Halo" ,"Vintage", "Other"])
metal_type = st.selectbox("Metal Type", ["Platinum", "Yellow Gold", "White Gold", "Rose Gold", "Other"])
stone_shape = st.selectbox("Center Stone Shape", ["Round", "Princess", "Oval", "Cushion", "Emerald", "Other"])
stone_type = st.selectbox("Stone Type", ["Diamond", "Sapphire", "Ruby", "Emerald", "Other"])
side_stones = st.text_input("Describe any side stones (or type 'none')", "none")
additional_details = st.text_input("Any additional details (e.g., engraving, setting style)")   

halo_stones = ""
if "halo" in ring_style.lower():
    halo_stones = st.text_input("Describe the halo stones")

button = st.button("Generate Images")

if button:
    base_desc = (
        f"An engagement ring with {ring_style.lower()} style, made of {metal_type.lower()}, "
        f"featuring a {stone_shape.lower()} {stone_type.lower()} center stone. Please keep in mind these additonal details when designing: {additional_details}" if additional_details else ""
    )
    if side_stones.lower() != "none":
        base_desc += f", with side stones: {side_stones}"
    if halo_stones:  # Removed .lower() check since halo_stones is already a string
        base_desc += f", with halo stones: {halo_stones}"

    views = {
        "top": f"{base_desc}, top view, studio lighting, crisp rendering, white background",
        "right": f"{base_desc}, right (profile) view, studio lighting, crisp rendering, white background",
        "back": f"{base_desc}, back view focusing on ring band, studio lighting, crisp rendering, white background",
        "perspective": f"{base_desc}, 3/4 perspective view, studio lighting, crisp rendering, white background",
    }

    # Generate and display each view
    for view, prompt in views.items():
        st.write(f"### {view.capitalize()} view")
        
        try:
            # Generate content for this specific view
            response = client.models.generate_content(
                model="gemini-2.0-flash-preview-image-generation",
                contents=prompt,
                config=types.GenerateContentConfig(response_modalities=["TEXT", "IMAGE"])
            )
            
            # Process the response for this view
            image_found = False
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'inline_data') and part.inline_data and part.inline_data.data:
                    image_bytes = part.inline_data.data
                    decoded_data = base64.b64decode(image_bytes)

                    try:
                        img = Image.open(BytesIO(decoded_data))
                        st.image(img, caption=f"{view.capitalize()} view", use_column_width=True)
                        image_found = True
                        break  # Only show the first image found for this view
                    except Exception as e:
                        st.error(f"Could not display image for {view} view. Error: {e}")
                        st.info(f"First 100 bytes: {image_bytes[:100]}")
            
            if not image_found:
                st.warning(f"No image generated for {view} view")
                
        except Exception as e:
            st.error(f"Error generating {view} view: {e}")
        
        # Add some spacing between views
        st.write("---")