Spaces:
Running
Running
Update app.py
Browse files
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 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
30 |
else:
|
31 |
-
|
32 |
-
|
33 |
-
|
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()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|