muneeb487 commited on
Commit
b887628
·
verified ·
1 Parent(s): 2ba9d62

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -44
app.py CHANGED
@@ -2,48 +2,48 @@ import streamlit as st
2
  from PIL import Image
3
  from diffusers import StableDiffusionInpaintPipeline
4
  import torch
5
- import os
6
-
7
- # Load Stable Diffusion Inpainting Pipeline
8
- @st.cache_resource
9
- def load_pipeline():
10
- pipeline = StableDiffusionInpaintPipeline.from_pretrained(
11
- "stabilityai/stable-diffusion-2-inpainting",
12
- torch_dtype=torch.float16,
13
- )
14
- return pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
15
-
16
- stability_pipeline = load_pipeline()
17
-
18
- # Streamlit App Title
19
- st.title("AI-Powered Image Editor")
20
- st.markdown("Upload an image, optionally upload a mask, and provide a command to edit the image.")
21
-
22
- # Image Upload
23
- uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
24
- uploaded_mask = st.file_uploader("Upload a Mask Image (optional)", type=["jpg", "jpeg", "png"])
25
- prompt = st.text_input("Enter your prompt (e.g., 'Add a sunset in the background')")
26
-
27
- if st.button("Generate"):
28
- if uploaded_image is None:
29
- st.error("Please upload an image to proceed.")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
30
  else:
31
- # Load and process the uploaded image
32
- image = Image.open(uploaded_image).convert("RGB")
33
- mask = Image.open(uploaded_mask).convert("RGB") if uploaded_mask else None
34
-
35
- st.image(image, caption="Uploaded Image", use_column_width=True)
36
- if mask:
37
- st.image(mask, caption="Uploaded Mask", use_column_width=True)
38
-
39
- # Generate the edited image
40
- with st.spinner("Generating edited image..."):
41
- try:
42
- result = stability_pipeline(prompt=prompt, image=image, mask_image=mask).images[0]
43
- st.image(result, caption="Edited Image", use_column_width=True)
44
- # Save the result
45
- output_path = "edited_image.jpg"
46
- result.save(output_path)
47
- st.success(f"Image generated and saved as {output_path}")
48
- except Exception as e:
49
- st.error(f"Error: {e}")
 
2
  from PIL import Image
3
  from diffusers import StableDiffusionInpaintPipeline
4
  import torch
5
+
6
+ # Load and display an image in Streamlit
7
+ def load_image(image_path):
8
+ # Open image
9
+ image = Image.open(image_path).convert('RGB')
10
+ return image
11
+
12
+ # Main function to process the image
13
+ def process_image(image, prompt):
14
+ # Create the pipeline using Stable Diffusion
15
+ pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting")
16
+
17
+ # If using GPU, send the pipeline to CUDA
18
+ pipe.to("cuda" if torch.cuda.is_available() else "cpu")
19
+
20
+ # Perform inpainting (change color and add fire in the background)
21
+ edited_image = pipe(prompt=prompt, init_image=image, strength=0.75).images[0]
22
+
23
+ return edited_image
24
+
25
+ # Streamlit Interface
26
+ def main():
27
+ # Upload image
28
+ uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
29
+
30
+ if uploaded_image is not None:
31
+ # Load image
32
+ image = load_image(uploaded_image)
33
+
34
+ # Display original image
35
+ st.image(image, caption="Original Image", use_container_width=True)
36
+
37
+ # Define the prompt
38
+ prompt = "change the color of dragon and add fire in the background"
39
+
40
+ # Process the image based on the prompt
41
+ edited_image = process_image(image, prompt)
42
+
43
+ # Display the edited image
44
+ st.image(edited_image, caption="Edited Image", use_container_width=True)
45
  else:
46
+ st.write("Please upload an image to begin.")
47
+
48
+ if __name__ == "__main__":
49
+ main()