Update app.py

app.py

@@ -1,66 +1,49 @@
 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
+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:
-                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)
+                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 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!
-""")
+                st.error(f"Error: {e}")