Spaces:
Running
Running
File size: 1,762 Bytes
786fe91 b303b1d dd27165 b887628 dd27165 95069ef b887628 dd27165 b887628 68e7afc 95069ef dd27165 95069ef b887628 dd27165 b887628 dd27165 b887628 dd27165 b887628 95069ef b887628 |
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
import torch
from diffusers import StableDiffusionInpaintPipeline
import numpy as np
# Load the StableDiffusionInpaintPipeline
pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting")
pipe.to("cuda" if torch.cuda.is_available() else "cpu") # Move model to GPU if available
# Function to process the image with the provided prompt
def process_image(image, prompt):
# Ensure the image is a PIL Image
if isinstance(image, np.ndarray):
image = Image.fromarray(image) # Convert numpy array to PIL Image
elif isinstance(image, torch.Tensor):
image = Image.fromarray(image.numpy()) # Convert torch tensor to PIL Image
# Resize image to the required size (512x512)
image = image.resize((512, 512))
# Process the image through the pipeline
edited_image = pipe(prompt=prompt, init_image=image, strength=0.75).images[0]
return edited_image
# Streamlit app
def main():
st.title("Image Inpainting with Stable Diffusion")
# Upload an image
uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
# Input prompt for the image editing
prompt = st.text_input("Enter your prompt", "Change the color of the dragon and add fire in the background")
if uploaded_file is not None and prompt:
# Open the uploaded image
image = Image.open(uploaded_file)
# Process the image based on the prompt
with st.spinner("Processing the image..."):
edited_image = process_image(image, prompt)
# Display the edited image
st.image(edited_image, caption="Edited Image", use_container_width=True)
if __name__ == "__main__":
main()
|