added inpainting module

#2
by Varhal - opened
Files changed (1) hide show
  1. inpaintModule.py +11 -0
inpaintModule.py ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import StableDiffusionInpaintPipeline
4
+ from PIL import Image
5
+
6
+ model_id = "stabilityai/stable-diffusion-2-inpainting"
7
+ pipeline = StableDiffusionInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
8
+
9
+ def inpaint_fn(image: Image.Image, mask: Image.Image, prompt: str):
10
+ result = pipeline(prompt=prompt, image=image, mask_image=mask).images[0]
11
+ return result