File size: 983 Bytes
3111b11 |
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 |
import gradio as gr
from PIL import Image
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image
def draw_on_image(image, prompt):
print(image, prompt)
if not prompt:
return
init_image = load_image(
image["image"]
)
mask_image = load_image(
image["mask"]
)
res_image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
return res_image
inputs = [
gr.Image(tool="sketch", label="Image", type="pil"),
gr.Text(max_lines=1)
]
if torch.cuda.is_available():
torch_dtype = torch.float32
device = "cuda"
else:
torch_dtype = torch.float16
device = "cpu"
pipeline = AutoPipelineForInpainting.from_pretrained(
"kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch_dtype
)
pipeline.to(device)
app = gr.Interface(draw_on_image, inputs=inputs, outputs="image")
app.queue()
app.launch(share=True) |