app.py

import streamlit as st
from PIL import Image
from diffusers import StableDiffusionInpaintPipeline
import torch
import os

# Load Stable Diffusion Inpainting Pipeline
@st.cache_resource
def load_pipeline():
    pipeline = StableDiffusionInpaintPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-inpainting",
        torch_dtype=torch.float16,
    )
    return pipeline.to("cuda" if torch.cuda.is_available() else "cpu")

stability_pipeline = load_pipeline()

# Streamlit App Title
st.title("AI-Powered Image Editor")
st.markdown("Upload an image, optionally upload a mask, and provide a command to edit the image.")

# Image Upload
uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
uploaded_mask = st.file_uploader("Upload a Mask Image (optional)", type=["jpg", "jpeg", "png"])
prompt = st.text_input("Enter your prompt (e.g., 'Add a sunset in the background')")

if st.button("Generate"):
    if uploaded_image is None:
        st.error("Please upload an image to proceed.")
    else:
        # Load and process the uploaded image
        image = Image.open(uploaded_image).convert("RGB")
        mask = Image.open(uploaded_mask).convert("RGB") if uploaded_mask else None

        st.image(image, caption="Uploaded Image", use_column_width=True)
        if mask:
            st.image(mask, caption="Uploaded Mask", use_column_width=True)

        # Generate the edited image
        with st.spinner("Generating edited image..."):
            try:
                result = stability_pipeline(prompt=prompt, image=image, mask_image=mask).images[0]
                st.image(result, caption="Edited Image", use_column_width=True)
                # Save the result
                output_path = "edited_image.jpg"
                result.save(output_path)
                st.success(f"Image generated and saved as {output_path}")
            except Exception as e:
                st.error(f"Error: {e}")