app.py

import streamlit as st
from openai import OpenAI
import requests
from PIL import Image
import io
import os
from datetime import datetime

def convert_to_png(image_file):
    """Convert uploaded image to PNG format and ensure it's under 4MB"""
    # Open the image using PIL
    image = Image.open(image_file)
    
    # Convert to RGBA if not already
    if image.mode in ('RGBA', 'RGB'):
        image = image.convert('RGBA')
    else:
        image = image.convert('RGB').convert('RGBA')
    
    # Save to bytes with PNG format
    byte_arr = io.BytesIO()
    image.save(byte_arr, format='PNG', optimize=True)
    byte_arr.seek(0)
    
    # Check if size is under 4MB
    if byte_arr.getbuffer().nbytes > 4 * 1024 * 1024:
        # If image is too large, resize it while maintaining aspect ratio
        while byte_arr.getbuffer().nbytes > 4 * 1024 * 1024:
            width, height = image.size
            new_width = int(width * 0.9)  # Reduce by 10%
            new_height = int(height * 0.9)
            image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
            
            byte_arr = io.BytesIO()
            image.save(byte_arr, format='PNG', optimize=True)
            byte_arr.seek(0)
    
    return byte_arr

def validate_image(uploaded_file):
    """Validate image size and format"""
    if uploaded_file.size > 4 * 1024 * 1024:  # 4MB in bytes
        return False, "File size must be less than 4MB"
    return True, "OK"

def save_uploaded_file(uploaded_file, folder="uploads"):
    """Save uploaded file to a temporary folder and return the path"""
    if not os.path.exists(folder):
        os.makedirs(folder)
    
    # Generate a timestamp for unique filename
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # Convert to PNG and validate size
    image_bytes = convert_to_png(uploaded_file)
    
    # Save the converted PNG
    file_path = os.path.join(folder, f"{timestamp}.png")
    with open(file_path, "wb") as f:
        f.write(image_bytes.getvalue())
    
    return file_path

def download_image(url, folder="generated"):
    """Download image from URL and save it"""
    if not os.path.exists(folder):
        os.makedirs(folder)
        
    response = requests.get(url)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    file_path = os.path.join(folder, f"generated_image_{timestamp}.png")
    
    if response.status_code == 200:
        with open(file_path, "wb") as f:
            f.write(response.content)
        return file_path
    return None

def main():
    st.title("🎨 Image Editor with DALL-E 2")
    
    # Sidebar for API key
    with st.sidebar:
        st.header("Configuration")
        api_key = st.text_input("Enter your OpenAI API key", type="password")
        st.markdown("""
        ### How to get an API key
        1. Go to [OpenAI API Keys](https://platform.openai.com/api-keys)
        2. Create a new secret key
        3. Copy and paste it here
        """)
        
        # Size selection
        size_option = st.selectbox(
            "Select image size:",
            ["1024x1024", "512x512", "256x256"]
        )
        
        # Number of images
        num_images = st.slider("Number of images to generate", 1, 4, 1)

    # Main content
    st.markdown("""
    ### Requirements:
    - Original image and mask must be less than 4MB
    - Images will be automatically converted to PNG format
    - Images larger than 4MB will be automatically resized
    """)

    # File uploaders
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Original Image")
        original_image = st.file_uploader(
            "Upload original image", 
            type=["png", "jpg", "jpeg"]
        )
        if original_image:
            # Validate image
            valid, message = validate_image(original_image)
            if not valid:
                st.warning(f"Original image: {message}. Image will be automatically resized.")
            
            try:
                # Display image preview
                image = Image.open(original_image)
                st.image(image, caption="Original Image", use_container_width=True)
                st.caption(f"Original size: {original_image.size/1024/1024:.2f}MB")
            except Exception as e:
                st.error(f"Error loading image: {str(e)}")

    with col2:
        st.subheader("Mask Image")
        mask_image = st.file_uploader(
            "Upload mask image", 
            type=["png", "jpg", "jpeg"]
        )
        if mask_image:
            # Validate image
            valid, message = validate_image(mask_image)
            if not valid:
                st.warning(f"Mask image: {message}. Image will be automatically resized.")
            
            try:
                # Display image preview
                image = Image.open(mask_image)
                st.image(image, caption="Mask Image", use_column_width=True)
                st.caption(f"Original size: {mask_image.size/1024/1024:.2f}MB")
            except Exception as e:
                st.error(f"Error loading mask: {str(e)}")

    # Prompt input
    prompt = st.text_area(
        "Enter your prompt:",
        placeholder="Describe the changes you want to make to the image...",
        help="Be specific about what you want to add or modify in the masked area"
    )

    # Generate button
    if st.button("Generate Edited Image"):
        if not api_key:
            st.error("Please enter your OpenAI API key in the sidebar.")
            return
            
        if not original_image or not mask_image:
            st.error("Please upload both an original image and a mask image.")
            return
            
        if not prompt:
            st.error("Please enter a prompt describing the desired changes.")
            return

        try:
            with st.spinner("Processing images and generating edited version..."):
                # Save and convert uploaded files
                original_path = save_uploaded_file(original_image)
                mask_path = save_uploaded_file(mask_image)

                # Initialize OpenAI client
                client = OpenAI(api_key=api_key)

                # Make the API call
                response = client.images.edit(
                    model="dall-e-2",
                    image=open(original_path, "rb"),
                    mask=open(mask_path, "rb"),
                    prompt=prompt,
                    n=num_images,
                    size=size_option
                )

                # Display results
                st.subheader("Generated Images")
                cols = st.columns(num_images)
                
                for idx, image_data in enumerate(response.data):
                    # Download and save the generated image
                    saved_image_path = download_image(image_data.url)
                    
                    if saved_image_path:
                        with cols[idx]:
                            st.image(saved_image_path, caption=f"Generated Image {idx+1}")
                            
                            # Add download button for each image
                            with open(saved_image_path, "rb") as file:
                                st.download_button(
                                    label=f"Download Image {idx+1}",
                                    data=file,
                                    file_name=f"edited_image_{idx+1}.png",
                                    mime="image/png"
                                )

                # Cleanup temporary files
                for path in [original_path, mask_path]:
                    if os.path.exists(path):
                        os.remove(path)

        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
            
        finally:
            # Cleanup temporary files
            for path in [original_path, mask_path]:
                if 'path' in locals() and os.path.exists(path):
                    try:
                        os.remove(path)
                    except Exception:
                        pass

if __name__ == "__main__":
    main()