Update inference.py

inference.py

@@ -1,7 +1,36 @@
+import streamlit as st
 import torch
 import numpy
 from PIL import Image
 from torchvision import transforms
 from diffusers import StableDiffusionInpaintPipeline
-from diffusers import DPMSolverMultistepScheduler
+from diffusers import DPMSolverMultistepScheduler, UniPCMultistepScheduler
 
+@torch.inference_mode()
+@st.experimental_singleton(max_entries=5)
+def get_pipeline():
+    pipe = StableDiffusionInpaintPipeline.from_pretrained("stabilityai/stable-diffusion-2-inpainting",
+                                                                      torch_dtype=torch.float16)
+    pipe.to(device)
+    pipe.enable_xformers_memory_efficient_attention()
+    pipe.set_progress_bar_config(disable=True)
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    return pipe
+
+
+def inpainting(image,
+               mask_image,
+               prompt,
+               negative_prompt,
+               num_inference_steps=20,
+               guidance_scale=7.5,
+               ):
+    pipe = get_pipeline()
+    result = pipe(
+        image=image,
+        mask_image=mask_image,
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+    ).images[0]