import streamlit as st
from PIL import Image, ImageDraw, ImageFont
import io
import requests
from transformers import pipeline

# Initialize the AI model for image generation or editing
# This uses HuggingFace's Stable Diffusion pipeline
stability_pipeline = pipeline("image-to-image", model="stabilityai/stable-diffusion-2-inpainting")

# Function to edit the image based on the command
def process_command(image, command):
    """
    Processes the command and modifies the image using AI tools.
    :param image: PIL Image object
    :param command: Natural language command (e.g., "Add a cat in the bottom-right corner.")
    :return: Edited PIL Image
    """
    # Convert the image to bytes for the AI model
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    img_byte_arr = img_byte_arr.getvalue()

    # Use the Stable Diffusion pipeline to modify the image
    result = stability_pipeline(images=image, prompt=command)
    edited_image = result[0]['image']
    return edited_image

# Streamlit App Interface
st.title("AI-Powered Image Editing")
st.write("Upload an image, enter a command, and let AI handle the editing!")

# Step 1: Upload Image
uploaded_file = st.file_uploader("Upload an image to edit", type=["jpg", "jpeg", "png"])

if uploaded_file:
    # Open the uploaded image
    input_image = Image.open(uploaded_file).convert("RGB")
    st.image(input_image, caption="Uploaded Image", use_column_width=True)

    # Step 2: Enter a Command
    command = st.text_input("Enter a command to edit the image (e.g., 'Add a tree in the background')")

    if st.button("Process Command"):
        if command:
            # Step 3: Process the Command
            try:
                with st.spinner("Processing your command..."):
                    edited_image = process_command(input_image, command)
                # Step 4: Display the Edited Image
                st.image(edited_image, caption="Edited Image", use_column_width=True)
            except Exception as e:
                st.error(f"Error processing the image: {e}")
        else:
            st.warning("Please enter a command.")

# Instructions for the user
st.write("""
### How to use this app:
1. Upload an image in JPG or PNG format.
2. Enter a natural language command like:
   - "Add a cat to the bottom-right corner."
   - "Remove the text from the top-left corner."
   - "Make the background cloudy."
3. Click "Process Command" and view the results!
""")