pretzinger commited on
Commit
ba89e30
1 Parent(s): 81a84f4

create app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -0
app.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from diffusers import StableDiffusionInpaintPipeline
3
+ from PIL import Image
4
+ import torch
5
+
6
+ # Set up title and description
7
+ st.title("Insert Yourself Into Historical Photos with Stable Diffusion!")
8
+ st.write("Upload a historical photo and a mask, then describe how you'd like to place yourself into the scene.")
9
+
10
+ # Load model (for inpainting)
11
+ @st.cache_resource
12
+ def load_pipeline():
13
+ pipe = StableDiffusionInpaintPipeline.from_pretrained(
14
+ "stabilityai/stable-diffusion-2-inpainting",
15
+ torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
16
+ )
17
+ pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
18
+ return pipe
19
+
20
+ pipe = load_pipeline()
21
+
22
+ # File uploads for base image and mask
23
+ base_image = st.file_uploader("Upload the Historical Photo", type=["jpg", "jpeg", "png"])
24
+ mask_image = st.file_uploader("Upload the Mask (Black & White Image)", type=["jpg", "jpeg", "png"])
25
+
26
+ # Prompt input
27
+ prompt = st.text_input("Describe the scene. How do you want to place yourself?")
28
+
29
+ if base_image and mask_image and prompt:
30
+ # Display input images
31
+ st.image(base_image, caption="Historical Photo", use_column_width=True)
32
+ st.image(mask_image, caption="Mask", use_column_width=True)
33
+
34
+ # Load images using PIL
35
+ base_image = Image.open(base_image).convert("RGB")
36
+ mask_image = Image.open(mask_image).convert("RGB")
37
+
38
+ # Generate inpainting
39
+ st.write("Generating the image... This may take a moment.")
40
+ result = pipe(prompt=prompt, image=base_image, mask_image=mask_image).images[0]
41
+
42
+ # Display result
43
+ st.image(result, caption="Inpainted Image", use_column_width=True)
44
+ st.write("Right-click on the image to download!")
45
+