Spaces:
Running
Running
File size: 1,946 Bytes
786fe91 b303b1d 786fe91 b303b1d 786fe91 b303b1d |
1 2 3 4 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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
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}")
|