Update app.py

app.py

@@ -2,48 +2,48 @@ 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.")
+
+# Load and display an image in Streamlit
+def load_image(image_path):
+    # Open image
+    image = Image.open(image_path).convert('RGB')
+    return image
+
+# Main function to process the image
+def process_image(image, prompt):
+    # Create the pipeline using Stable Diffusion
+    pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting")
+    
+    # If using GPU, send the pipeline to CUDA
+    pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Perform inpainting (change color and add fire in the background)
+    edited_image = pipe(prompt=prompt, init_image=image, strength=0.75).images[0]
+    
+    return edited_image
+
+# Streamlit Interface
+def main():
+    # Upload image
+    uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+    
+    if uploaded_image is not None:
+        # Load image
+        image = load_image(uploaded_image)
+
+        # Display original image
+        st.image(image, caption="Original Image", use_container_width=True)
+
+        # Define the prompt
+        prompt = "change the color of dragon and add fire in the background"
+        
+        # Process the image based on the prompt
+        edited_image = process_image(image, prompt)
+
+        # Display the edited image
+        st.image(edited_image, caption="Edited Image", use_container_width=True)
     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}")
+        st.write("Please upload an image to begin.")
+
+if __name__ == "__main__":
+    main()